Align DreamWaQ PPO rollout and policy dynamics

motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py

@@ -197,7 +197,7 @@ class DreamWaQTask(Go1WalkTask):
 
     def apply_action(self, actions, state):
         """裁剪动作 + 随机延迟（模拟真实部署延迟，提高 sim-to-real 泛化）。"""
-        actions = np.clip(actions, -4.0, 4.0)
+        actions = np.clip(actions, -100.0, 100.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)

motrix_rl/src/motrix_rl/rslrl/cfg.py

@@ -36,7 +36,7 @@ class RslRlActorCfg:
     class_name: str = "MLPModel"
     hidden_dims: list[int] = field(default_factory=lambda: [256, 128, 64])
     activation: str = "elu"
-    obs_normalization: bool = True
+    obs_normalization: bool = False
     stochastic: bool = True
     init_noise_std: float = 1.0
     noise_std_type: Literal["scalar", "log"] = "scalar"

motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py

@@ -114,8 +114,7 @@ class CENetActorModel(MLPModel):
 
         # 在 nn.Module.__init__ 之后创建 VAE 子模块
         self.vae = CENetVAE(cenet_in_dim, cenet_out_dim, activation)
-        self.action_clip = 4.0
+        self.std_floor = 1e-6
-        self.std_clip = 0.6
         self._last_cenet_output = None
         # AdaBoot: 自适应 bootstrapping（论文 Section II-C）
         self._adaboot_cv_buffer = []  # 速度估计误差的 CV 历史
@@ -123,15 +122,16 @@ class CENetActorModel(MLPModel):
 
     def _update_distribution(self, obs: torch.Tensor) -> None:
         """覆盖父类 — 限制动作分布，避免采样动作进入 PD/关节限位饱和区。"""
-        mean = torch.clamp(self.mlp(obs), -self.action_clip, self.action_clip)
+        mean = self.mlp(obs)
         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.nan_to_num_(nan=1.0, posinf=1.0, neginf=1.0)
-                    self.std.clamp_(min=1e-6, max=self.std_clip)
+                    self.std.clamp_(min=self.std_floor)
                 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=torch.log(torch.tensor(self.std_clip)).item())
+                    self.log_std.nan_to_num_(nan=0.0, posinf=5.0, neginf=-20.0)
+                    self.log_std.clamp_(min=-20.0)
             if self.noise_std_type == "scalar":
                 std = self.std.expand_as(mean)
             elif self.noise_std_type == "log":
@@ -139,7 +139,7 @@ class CENetActorModel(MLPModel):
             else:
                 raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}.")
         else:
-            std = torch.full_like(mean, self.std_clip)
+            std = torch.ones_like(mean)
         self.distribution = Normal(mean, std)
 
     def forward(
@@ -150,7 +150,7 @@ class CENetActorModel(MLPModel):
         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)
+        return actions
 
     def _get_latent_dim(self) -> int:
         """Actor 实际输入：code(19) + policy(45) = 64。

motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py

@@ -85,6 +85,7 @@ class DreamWaQPPO(PPO):
             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 ("prev_privileged_obs" in obs_batch and not torch.isfinite(obs_batch["prev_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()
@@ -292,7 +293,7 @@ class DreamWaQPPO(PPO):
         code, code_vel, decode, mean_vel, logvar_vel, mean_latent, logvar_latent = cenet_out
 
         # 速度估计目标：privileged_obs 中的 base_vel（索引 45:48）
-        vel_target = obs_batch["privileged_obs"][:, 45:48]
+        vel_target = obs_batch["prev_privileged_obs"][:, 45:48]
         # 观测重建目标：当前 policy obs
         obs_target = obs_batch["policy"]
 

motrix_rl/src/motrix_rl/rslrl/torch/wrap_vec_env.py

@@ -83,7 +83,7 @@ class RslrlNpEnvWrap(VecEnv):
         """Return the unwrapped environment (self for this wrapper)."""
         return self
 
-    def _build_obs_dict(self, state) -> dict[str, torch.Tensor]:
+    def _build_obs_dict(self, state, prev_privileged_obs=None) -> dict[str, torch.Tensor]:
         """将 NpEnvState 中的观测字段组装为 TensorDict 字典。
 
         支持 env 通过 state.info 传递 obs_history 和 privileged_obs。
@@ -96,12 +96,22 @@ class RslrlNpEnvWrap(VecEnv):
         if "privileged_obs" in state.info:
             priv = np.nan_to_num(state.info["privileged_obs"], nan=0.0)
             obs_dict["privileged_obs"] = torch.from_numpy(priv).to(self._device)
+            if prev_privileged_obs is None:
+                prev_privileged_obs = np.zeros_like(priv)
+            obs_dict["prev_privileged_obs"] = torch.from_numpy(
+                np.nan_to_num(prev_privileged_obs, nan=0.0)
+            ).to(self._device)
         return obs_dict
 
     def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]:
         # Convert torch actions to numpy
         actions_np = actions.cpu().numpy()
 
+        # Preserve the privileged observation aligned with the policy history.
+        previous_privileged_obs = None
+        if self._state is not None and "privileged_obs" in self._state.info:
+            previous_privileged_obs = self._state.info["privileged_obs"].copy()
+
         # Step the environment
         state = self._env.step(actions_np)
         self._state = state
@@ -119,7 +129,7 @@ class RslrlNpEnvWrap(VecEnv):
         dones = torch.from_numpy(state.done.astype(np.float32)).to(self._device)
 
         # 构建多键 TensorDict（policy + obs_history + privileged_obs）
-        obs = TensorDict(self._build_obs_dict(state),
+        obs = TensorDict(self._build_obs_dict(state, previous_privileged_obs),
                          batch_size=[self._num_envs], device=self._device)
 
         # Build extras dict (RSLRL calls it "extras" not "infos")

motrix_rl/src/motrix_rl/tasks/go1.py

@@ -126,7 +126,7 @@ class rslrl:
             runner.algorithm.entropy_coef = 0.01      # 上游 Go1RoughCfgPPO
             runner.algorithm.desired_kl = 0.01        # 上游 0.01 (默认 0.008)
             runner.algorithm.clip_param = 0.2
-            runner.algorithm.schedule = "fixed"  # 4096 envs 采样够大，不需要 adaptive 加噪
+            runner.algorithm.schedule = "adaptive"
             runner.algorithm.gamma = 0.99
             runner.algorithm.lam = 0.95
             runner.algorithm.max_grad_norm = 1.0

scripts/dreamwaq_sim2sim_mujoco.py

@@ -39,7 +39,7 @@ HISTORY_LEN = 5
 ACTION_SCALE = 0.25
 KP = 28.0
 KD = 0.7
-CLIP_ACTIONS = 4.0
+CLIP_ACTIONS = 100.0
 CLIP_TORQUES = 80.0
 CLIP_OBS = 100.0
 MAX_VX, MAX_VY, MAX_WZ = 1.0, 1.0, 1.0