feat: DreamWaQ full replication — env, terrain, CENet, PPO

motrix_rl/src/motrix_rl/rslrl/cfg.py

@@ -74,6 +74,10 @@ class RslRlPpoAlgorithmCfg:
     max_grad_norm: float = 1.0
     normalize_advantage_per_mini_batch: bool = False
     rnd_cfg: dict | None = None
+    # DreamWaQ CENet 参数
+    vae_beta: float = 1.0       # VAE KL 散度权重
+    cenet_in_dim: int = 225     # 观测历史维度 (num_history × obs_dim)
+    cenet_out_dim: int = 19     # CENet code 维度 (vel_est 3 + latent 16)
     symmetry_cfg: dict | None = None
 
 

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

@@ -0,0 +1,2 @@
+# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
+"""RSLRL 自定义模型模块。"""

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

@@ -0,0 +1,162 @@
+# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
+"""CENet Actor 模型——DreamWaQ 的 VAE 编码器 + MLP Actor。
+
+架构：
+  CENetVAE:  obs_history(225) → encoder[128,64] → latent(16) + vel_est(3) = code(19)
+             code(19) → decoder[64,128] → 重建 obs(45)
+  CENetActorModel (继承 MLPModel):
+              code(19) + obs(45) = 64 → MLP[512,256,128] → action(12)
+
+用法：
+  class_name = "motrix_rl.rslrl.torch.models.cenet_actor:CENetActorModel"
+  obs_groups = {"actor": ["policy", "obs_history"], "critic": ["privileged_obs"]}
+"""
+
+import torch
+import torch.nn as nn
+from tensordict import TensorDict
+
+from rsl_rl.models.mlp_model import MLPModel
+from rsl_rl.modules import EmpiricalNormalization, HiddenState
+
+
+# ═══ CENet VAE ═══
+
+class CENetVAE(nn.Module):
+    """CENet VAE 编码器-解码器。
+
+    obs_history(225) → encoder → latent(16) + vel_est(3) = code(19)
+    code(19) → decoder → 重建 obs(45)
+    """
+
+    def __init__(self, cenet_in_dim: int = 225, cenet_out_dim: int = 19,
+                 activation: str = "elu"):
+        super().__init__()
+        act = _get_activation(activation)
+        self.cenet_in_dim = cenet_in_dim
+        self.cenet_out_dim = cenet_out_dim
+
+        # 编码器：225 → 128 → 64
+        self.encoder = nn.Sequential(
+            nn.Linear(cenet_in_dim, 128), act,
+            nn.Linear(128, 64), act,
+        )
+        # 潜变量头：64 → 16 (mean + logvar)
+        self.encode_mean_latent = nn.Linear(64, cenet_out_dim - 3)
+        self.encode_logvar_latent = nn.Linear(64, cenet_out_dim - 3)
+        # 速度估计头：64 → 3 (mean + logvar)
+        self.encode_mean_vel = nn.Linear(64, 3)
+        self.encode_logvar_vel = nn.Linear(64, 3)
+        # 解码器：19 → 64 → 128 → 45
+        self.decoder = nn.Sequential(
+            nn.Linear(cenet_out_dim, 64), act,
+            nn.Linear(64, 128), act,
+            nn.Linear(128, 45),
+        )
+
+    def reparameterise(self, mean: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+        """重参数化技巧：从 N(mean, exp(logvar/2)) 采样。"""
+        std = torch.exp(logvar * 0.5)
+        eps = torch.randn_like(std)
+        return mean + std * eps
+
+    def forward(self, obs_history: torch.Tensor):
+        """前向传播。
+
+        Returns:
+            code:       (N, 19)  潜变量 [vel_sample(3) + latent_sample(16)]
+            code_vel:   (N, 3)   速度估计采样
+            decode:     (N, 45)  重建观测
+            mean_vel:   (N, 3)   速度估计均值
+            logvar_vel: (N, 3)   速度估计对数方差
+            mean_latent:   (N, 16)  潜变量均值
+            logvar_latent: (N, 16)  潜变量对数方差
+        """
+        h = self.encoder(obs_history)
+        mean_latent = self.encode_mean_latent(h)
+        logvar_latent = self.encode_logvar_latent(h)
+        mean_vel = self.encode_mean_vel(h)
+        logvar_vel = self.encode_logvar_vel(h)
+        code_latent = self.reparameterise(mean_latent, logvar_latent)
+        code_vel = self.reparameterise(mean_vel, logvar_vel)
+        code = torch.cat((code_vel, code_latent), dim=-1)
+        decode = self.decoder(code)
+        return code, code_vel, decode, mean_vel, logvar_vel, mean_latent, logvar_latent
+
+    def deterministic_code(self, obs_history: torch.Tensor) -> torch.Tensor:
+        """推理模式：使用均值而非采样，产生确定性 code(19)。"""
+        h = self.encoder(obs_history)
+        mean_latent = self.encode_mean_latent(h)
+        mean_vel = self.encode_mean_vel(h)
+        return torch.cat((mean_vel, mean_latent), dim=-1)
+
+
+# ═══ CENet Actor Model ═══
+
+class CENetActorModel(MLPModel):
+    """CENet Actor：继承 MLPModel，在 get_latent() 中注入 VAE code。
+
+    class_name = "motrix_rl.rslrl.torch.models.cenet_actor:CENetActorModel"
+    """
+
+    def __init__(self, obs: TensorDict, obs_groups: dict[str, list[str]],
+                 obs_set: str, output_dim: int,
+                 cenet_in_dim: int = 225, cenet_out_dim: int = 19,
+                 activation: str = "elu", **kwargs):
+        # 必须在 super().__init__ 之前设置，因为 _get_latent_dim() 会被父类构造函数调用
+        self._history_dim = cenet_in_dim
+        self._code_dim = cenet_out_dim
+        # 禁用观测归一化（VAE 输出已是归一化后的 code，维度也不匹配）
+        kwargs["obs_normalization"] = False
+
+        super().__init__(obs, obs_groups, obs_set, output_dim, **kwargs)
+
+        # 在 nn.Module.__init__ 之后创建 VAE 子模块
+        self.vae = CENetVAE(cenet_in_dim, cenet_out_dim, activation)
+        self._last_cenet_output = None
+
+    def _update_distribution(self, obs: torch.Tensor) -> None:
+        """覆盖父类 — 确保 std 始终为正，防止 NaN。"""
+        super()._update_distribution(obs)
+        # 如果 std 因数值问题变负，clamp 到最小值
+        if self.stochastic and not self.state_dependent_std:
+            with torch.no_grad():
+                if self.noise_std_type == "scalar":
+                    self.std.clamp_(min=1e-6)
+                elif self.noise_std_type == "log":
+                    self.log_std.clamp_(min=-20.0, max=10.0)
+
+    def _get_latent_dim(self) -> int:
+        """Actor 实际输入：code(19) + policy(45) = 64。"""
+        return self.obs_dim - self._history_dim + self._code_dim
+
+    def get_latent(self, obs: TensorDict, masks: torch.Tensor | None = None,
+                   hidden_state: HiddenState = None) -> torch.Tensor:
+        """提取观测 → VAE 编码 → 拼接 code + policy → 返回 latent(64)。"""
+        policy_obs = obs["policy"]           # (N, 45)
+        obs_history = obs["obs_history"]     # (N, 225)
+
+        out = self.vae(obs_history)
+        self._last_cenet_output = out
+        code, code_vel, decode, mean_vel, logvar_vel, mean_latent, logvar_latent = out
+
+        latent = torch.cat([code, policy_obs], dim=-1)  # (N, 64)
+        return latent
+
+    def update_normalization(self, obs: TensorDict) -> None:
+        """CENetActor 使用 obs_normalization=False，此方法为空。"""
+        pass
+
+
+# ═══ 工具函数 ═══
+
+def _get_activation(act_name: str) -> nn.Module:
+    """解析激活函数名称。"""
+    _map = {
+        "elu": nn.ELU, "selu": nn.SELU, "relu": nn.ReLU,
+        "lrelu": nn.LeakyReLU, "tanh": nn.Tanh, "sigmoid": nn.Sigmoid,
+    }
+    if act_name in _map:
+        return _map[act_name]()
+
+    raise ValueError(f"未知激活函数: {act_name}。可选: {list(_map.keys())}")

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

@@ -0,0 +1,205 @@
+# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
+"""DreamWaQ PPO——标准 PPO + CENet VAE 自编码器损失。
+
+在 PPO.update() 的每个 mini-batch 中，额外计算：
+  - 速度估计损失：MSE(code_vel, 真实 base_vel)（来自 privileged_obs）
+  - 观测重建损失：MSE(decode, 当前 policy obs)
+  - KL 散度损失：beta * KL(N(mean, var) || N(0, 1))
+
+class_name = "motrix_rl.rslrl.torch.train.dreamwaq_ppo:DreamWaQPPO"
+"""
+
+import torch
+import torch.nn as nn
+from tensordict import TensorDict
+
+from rsl_rl.algorithms import PPO
+from rsl_rl.storage import RolloutStorage
+from rsl_rl.utils import resolve_callable, resolve_obs_groups
+
+
+class DreamWaQPPO(PPO):
+    """PPO + Beta-VAE 自编码器损失——DreamWaQ CENet 训练。"""
+
+    # VAE 损失权重（与上游 beta=1.0 一致）
+    vae_beta: float = 1.0
+
+    @staticmethod
+    def construct_algorithm(obs: TensorDict, env, cfg: dict, device: str) -> "DreamWaQPPO":
+        """构造 DreamWaQ PPO 算法——创建 CENetActor + MLP Critic。
+
+        与父类 PPO.construct_algorithm 的区别：
+          - actor 使用 CENetActorModel（含 VAE）
+          - obs_groups 中 actor=["policy", "obs_history"], critic=["privileged_obs"]
+        """
+        # 提取 DreamWaQ 特有参数
+        vae_beta = cfg.pop("vae_beta", 1.0)
+
+        # 解析 actor / critic 类
+        actor_class = resolve_callable(cfg["actor"].pop("class_name"))
+        critic_class = resolve_callable(cfg["critic"].pop("class_name"))
+
+        # 解析观测分组
+        obs_groups = resolve_obs_groups(obs, cfg["obs_groups"], ["actor", "critic"])
+
+        # 创建 actor（CENetActorModel）
+        actor = actor_class(obs, obs_groups, "actor", env.num_actions, **cfg["actor"]).to(device)
+
+        # 创建 critic（标准 MLPModel，输入 privileged_obs）
+        critic = critic_class(obs, obs_groups, "critic", 1, **cfg["critic"]).to(device)
+
+        # 初始化 rollout 存储
+        storage = RolloutStorage(
+            "rl", env.num_envs, cfg["num_steps_per_env"], obs, [env.num_actions], device
+        )
+
+        # 提取算法参数（移除 DreamWaQPPO 特有 key，剩余传给父类 PPO.__init__）
+        algo_cfg = dict(cfg["algorithm"])
+        for dw_key in ("class_name", "vae_beta", "cenet_in_dim", "cenet_out_dim"):
+            algo_cfg.pop(dw_key, None)
+
+        # 创建 DreamWaQPPO 实例
+        alg = DreamWaQPPO(actor, critic, storage, device=device, **algo_cfg)
+        alg.vae_beta = vae_beta
+        return alg
+
+    def update(self) -> dict[str, float]:
+        """标准 PPO update + VAE 自编码器损失。"""
+        mean_value_loss = 0.0
+        mean_surrogate_loss = 0.0
+        mean_autoenc_loss = 0.0
+
+        if self.actor.is_recurrent or self.critic.is_recurrent:
+            generator = self.storage.recurrent_mini_batch_generator(
+                self.num_mini_batches, self.num_learning_epochs)
+        else:
+            generator = self.storage.mini_batch_generator(
+                self.num_mini_batches, self.num_learning_epochs)
+
+        for (
+            obs_batch, actions_batch, target_values_batch, advantages_batch,
+            returns_batch, old_actions_log_prob_batch, old_mu_batch,
+            old_sigma_batch, hid_states_batch, masks_batch,
+        ) in generator:
+            # ── 标准 PPO 前向 ──
+            self.actor(obs_batch, masks=masks_batch, stochastic_output=True)
+            self.critic(obs_batch, masks=masks_batch)
+
+            # ── 动作对数概率 ──
+            actions_log_prob_batch = self.actor.get_output_log_prob(actions_batch)
+
+            # ── 价值损失 ──
+            value_batch = target_values_batch
+            if self.use_clipped_value_loss:
+                value_clipped = target_values_batch + (self.critic(obs_batch).detach() - target_values_batch).clamp(
+                    -self.clip_param, self.clip_param)
+                value_losses = torch.square(self.critic(obs_batch) - returns_batch)
+                value_losses_clipped = torch.square(value_clipped - returns_batch)
+                value_loss = torch.max(value_losses, value_losses_clipped).mean()
+            else:
+                value_loss = torch.square(returns_batch - self.critic(obs_batch)).mean()
+
+            # ── 代理损失 ──
+            surrogate_loss = self._compute_surrogate_loss(
+                actions_batch, actions_log_prob_batch,
+                old_actions_log_prob_batch, advantages_batch)
+
+            # ── 熵 ──
+            entropy_batch = self.actor.output_entropy
+            entropy_loss = entropy_batch.mean()
+
+            # ── VAE 自编码器损失 ──
+            autoenc_loss = self._compute_vae_loss(obs_batch)
+
+            # ── 总损失 ──
+            loss = (
+                surrogate_loss
+                + self.value_loss_coef * value_loss
+                - self.entropy_coef * entropy_loss
+                + autoenc_loss
+            )
+
+            # ── 梯度更新 ──
+            self.optimizer.zero_grad()
+            loss.backward()
+            if self.max_grad_norm is not None:
+                nn.utils.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm)
+                nn.utils.clip_grad_norm_(self.critic.parameters(), self.max_grad_norm)
+                nn.utils.clip_grad_value_(self.actor.parameters(), 10.0)
+            self.optimizer.step()
+            # 每次更新后强制 std > 0，防止数值异常导致 NaN
+            if hasattr(self.actor, 'std') and self.actor.stochastic:
+                with torch.no_grad():
+                    self.actor.std.clamp_(min=1e-6)
+
+            # ── 累计日志 ──
+            mean_value_loss += value_loss.item()
+            mean_surrogate_loss += surrogate_loss.item()
+            mean_autoenc_loss += autoenc_loss.item()
+
+        num_updates = self.num_learning_epochs * self.num_mini_batches
+        mean_value_loss /= num_updates
+        mean_surrogate_loss /= num_updates
+        mean_autoenc_loss /= num_updates
+
+        self.storage.clear()
+
+        return {
+            "value_loss": mean_value_loss,
+            "surrogate_loss": mean_surrogate_loss,
+            "autoenc_loss": mean_autoenc_loss,
+        }
+
+    def _compute_surrogate_loss(
+        self, actions_batch, actions_log_prob_batch,
+        old_actions_log_prob_batch, advantages_batch
+    ) -> torch.Tensor:
+        """计算 PPO 代理损失（从父类 PPO.update() 中提取）。"""
+        ratio = torch.exp(actions_log_prob_batch - old_actions_log_prob_batch)
+        surrogate = -advantages_batch * ratio
+        surrogate_clipped = -advantages_batch * torch.clamp(
+            ratio, 1.0 - self.clip_param, 1.0 + self.clip_param)
+        return torch.max(surrogate, surrogate_clipped).mean()
+
+    def _compute_vae_loss(self, obs_batch: TensorDict) -> torch.Tensor:
+        """计算 CENet VAE 损失。
+
+        obs_batch 包含：
+          - "policy": 当前观测 (N, 45) = 重建目标
+          - "obs_history": 观测历史 (N, 225) = VAE 编码器输入
+          - "privileged_obs": 特权观测 (N, 247), 其中 [45:48] 是 base_vel
+
+        损失组成：
+          1. 速度估计损失：MSE(code_vel, base_vel_gt)
+          2. 观测重建损失：MSE(decoded_obs, policy_obs)
+          3. KL 散度：beta * KL(q(z|history) || N(0,1))
+        """
+        # 从 actor 获取最近一次 CENet 前向输出
+        cenet_out = getattr(self.actor, "_last_cenet_output", None)
+        if cenet_out is None:
+            return torch.tensor(0.0, device=obs_batch.device)
+
+        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]
+        # 观测重建目标：当前 policy obs
+        obs_target = obs_batch["policy"]
+
+        mse = nn.functional.mse_loss
+        estimation_loss = mse(code_vel, vel_target)
+        reconstruction_loss = mse(decode, obs_target)
+        # KL 散度：-0.5 * sum(1 + logvar - mean^2 - exp(logvar))
+        # clamp logvar 防止 exp 溢出
+        logvar_latent = torch.clamp(logvar_latent, -20.0, 10.0)
+        kl_loss = -0.5 * torch.sum(
+            1 + logvar_latent - mean_latent.pow(2) - logvar_latent.exp(), dim=-1
+        ).mean()
+
+        autoenc_loss = (
+            estimation_loss + reconstruction_loss + self.vae_beta * kl_loss
+        )
+        # 防止 NaN 传播
+        if torch.isnan(autoenc_loss) or torch.isinf(autoenc_loss):
+            return torch.tensor(0.0, device=obs_batch.device)
+        return autoenc_loss

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

@@ -16,6 +16,7 @@
 """PPO Trainer for RSLRL integration."""
 
 import logging
+import os
 
 import torch
 from rsl_rl.runners import OnPolicyRunner
@@ -48,6 +49,7 @@ class Trainer:
         sim_backend: str = None,
         enable_render: bool = False,
         cfg_override: dict = None,
+        env_cfg_override: dict = None,
     ) -> None:
         """Initialize the RSLRL PPO trainer.
 
@@ -56,6 +58,7 @@ class Trainer:
             sim_backend: Simulation backend to use (e.g., "mujoco", "npcm")
             enable_render: Whether to enable rendering during training
             cfg_override: Optional configuration overrides
+            env_cfg_override: Optional env config overrides passed to make()
         """
         rlcfg = rl_registry.default_rl_cfg(env_name, "rslrl", backend="torch")
         if cfg_override is not None:
@@ -64,23 +67,29 @@ class Trainer:
         self._env_name = env_name
         self._sim_backend = sim_backend
         self._enable_render = enable_render
+        self._env_cfg_override = env_cfg_override
 
-    def train(self) -> None:
+    def train(self, checkpoint: str = None) -> None:
         """Start training the agent.
 
         Creates the environment, wraps it for RSLRL, and runs the training loop.
+
+        Args:
+            checkpoint: Optional path to a checkpoint (.pt) to resume from.
         """
         rlcfg = self._rlcfg
 
         # Create environment
-        env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.num_envs)
+        env = env_registry.make(self._env_name, sim_backend=self._sim_backend,
+                                num_envs=rlcfg.num_envs, env_cfg_override=self._env_cfg_override)
 
         # Set random seed
         if rlcfg.runner.seed is not None:
             torch.manual_seed(rlcfg.runner.seed)
 
-        # Determine device
-        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        # Determine device（可通过 MOTRIX_DEVICE=cpu 强制 CPU 训练）
+        device_str = os.environ.get("MOTRIX_DEVICE", "cuda:0" if torch.cuda.is_available() else "cpu")
+        device = torch.device(device_str)
         logger.info(f"Using device: {device}")
 
         # Wrap environment for RSLRL
@@ -94,6 +103,11 @@ class Trainer:
             vec_env, rslrl_cfg, log_dir=get_log_dir(self._env_name, rllib="rslrl", agent_name="PPO"), device=device
         )
 
+        # Load checkpoint if specified
+        if checkpoint:
+            runner.load(checkpoint)
+            logger.info(f"Resumed from checkpoint: {checkpoint}")
+
         # Start training
         logger.info(f"Starting training for {self._env_name}")
         logger.info(f"Number of environments: {rlcfg.num_envs}")
@@ -117,14 +131,25 @@ class Trainer:
         rlcfg = self._rlcfg
 
         # Create environment with play_num_envs
-        env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
+        # Enable play_mode if the env config supports it (for multi-terrain random spawn)
+        play_override = {"play_mode": True}
+        if self._env_cfg_override:
+            play_override.update(self._env_cfg_override)
+        try:
+            env = env_registry.make(
+                self._env_name, sim_backend=self._sim_backend,
+                num_envs=rlcfg.play_num_envs, env_cfg_override=play_override,
+            )
+        except ValueError:
+            env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
 
         # Set random seed
         if rlcfg.runner.seed is not None:
             torch.manual_seed(rlcfg.runner.seed)
 
         # Determine device
-        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        device_str = os.environ.get("MOTRIX_DEVICE", "cuda:0" if torch.cuda.is_available() else "cpu")
+        device = torch.device(device_str)
 
         # Wrap environment for RSLRL
         vec_env = RslrlNpEnvWrap(env, device)

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

@@ -83,6 +83,21 @@ class RslrlNpEnvWrap(VecEnv):
         """Return the unwrapped environment (self for this wrapper)."""
         return self
 
+    def _build_obs_dict(self, state) -> dict[str, torch.Tensor]:
+        """将 NpEnvState 中的观测字段组装为 TensorDict 字典。
+
+        支持 env 通过 state.info 传递 obs_history 和 privileged_obs。
+        """
+        obs_dict = {"policy": torch.from_numpy(state.obs).to(self._device)}
+        if "obs_history" in state.info:
+            hist = state.info["obs_history"]  # (N, num_history, obs_dim)
+            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)
+        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()
@@ -98,20 +113,24 @@ class RslrlNpEnvWrap(VecEnv):
         self.episode_length_buf[dones_np] = 0
 
         # Convert to torch tensors
-        obs_tensor = torch.from_numpy(state.obs).to(self._device)
         rewards = torch.from_numpy(state.reward).to(self._device)
 
         # Merge terminated and truncated into dones
         dones = torch.from_numpy(state.done.astype(np.float32)).to(self._device)
 
-        # Create TensorDict for observations
-        obs = TensorDict({"policy": obs_tensor}, batch_size=[self._num_envs], device=self._device)
+        # 构建多键 TensorDict（policy + obs_history + privileged_obs）
+        obs = TensorDict(self._build_obs_dict(state),
+                         batch_size=[self._num_envs], device=self._device)
 
         # Build extras dict (RSLRL calls it "extras" not "infos")
         extras = {}
         if "time_outs" in state.info:
             extras["time_outs"] = torch.from_numpy(state.info["time_outs"]).to(self._device)
 
+        # 将 episode 各项奖励传入 TensorBoard
+        if "ep_report" in state.info:
+            extras["episode"] = state.info["ep_report"]
+
         return obs, rewards, dones, extras
 
     def reset(self) -> tuple[TensorDict, dict]:
@@ -128,10 +147,9 @@ class RslrlNpEnvWrap(VecEnv):
         # Reset episode length buffer
         self.episode_length_buf.zero_()
 
-        obs_tensor = torch.from_numpy(state.obs).to(self._device)
-
-        # Create TensorDict for observations
-        obs = TensorDict({"policy": obs_tensor}, batch_size=[self._num_envs], device=self._device)
+        # 构建多键 TensorDict
+        obs = TensorDict(self._build_obs_dict(state),
+                         batch_size=[self._num_envs], device=self._device)
 
         # Build extras dict
         extras = {}
@@ -139,17 +157,17 @@ class RslrlNpEnvWrap(VecEnv):
         return obs, extras
 
     def get_observations(self) -> TensorDict:
-        """Get current observations without stepping the environment.
+        """获取当前观测（不步进环境）。
 
         Returns:
-            Current observations as TensorDict
+            当前观测的 TensorDict（含 policy, obs_history, privileged_obs）
         """
         if self._state is None:
             obs, _ = self.reset()
             return obs
 
-        obs_tensor = torch.from_numpy(self._state.obs).to(self._device)
-        obs = TensorDict({"policy": obs_tensor}, batch_size=[self._num_envs], device=self._device)
+        obs = TensorDict(self._build_obs_dict(self._state),
+                         batch_size=[self._num_envs], device=self._device)
         return obs
 
     def render(self) -> None:

motrix_rl/src/motrix_rl/tasks/go1.py

@@ -66,6 +66,19 @@ class skrl:
     @dataclass
     class Go1WalkStairsPPO(Go1WalkRoughSkrlPpo): ...
 
+    @rlcfg("go1-stairs-terrain-walk-no-linevel")
+    @dataclass
+    class Go1WalkStairsNoLinvelSkrlPpo(Go1WalkRoughSkrlPpo):
+        """Go1 stairs terrain walk (no linear velocity obs) - SKRL PPO config.
+
+        Uses [512, 256, 128] network from rough terrain config.
+        Increased timesteps for 4-phase curriculum.
+        """
+
+        def __post_init__(self):
+            super().__post_init__()
+            self.runner.trainer.timesteps = 60000
+
 
 class rslrl:
     @rlcfg("go1-flat-terrain-walk")
@@ -93,6 +106,53 @@ class rslrl:
             algo.num_learning_epochs = 5
             algo.num_mini_batches = 3
 
+    @rlcfg("go1-dreamwaq-walk")
+    @dataclass
+    class Go1DreamWaQWalkRslrlPpo(RslrlCfg):
+        """Go1 DreamWaQ walk — CENet VAE + 不对称特权观测。"""
+
+        num_envs: int = 1024  # 上游 4096，CPU/GPU 安全默认
+
+        def __post_init__(self):
+            runner = self.runner
+
+            # Runner 设置（严格对齐上游 LeggedRobotCfgPPO + Go1RoughCfgPPO）
+            runner.seed = 5                     # 上游 seed=5
+            runner.max_iterations = 3000
+            runner.num_steps_per_env = 24
+            runner.experiment_name = "go1_dreamwaq_walk"
+            runner.save_interval = 50
+
+            # 算法：DreamWaQPPO（含 VAE loss）—— 严格对齐上游
+            runner.algorithm.class_name = (
+                "motrix_rl.rslrl.torch.train.dreamwaq_ppo:DreamWaQPPO")
+            runner.algorithm.learning_rate = 1e-3     # 上游 1.e-3
+            runner.algorithm.num_learning_epochs = 5
+            runner.algorithm.num_mini_batches = 4
+            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.gamma = 0.99
+            runner.algorithm.lam = 0.95
+            runner.algorithm.max_grad_norm = 1.0
+            runner.algorithm.vae_beta = 1.0
+
+            # Actor：CENetActorModel（code 替换 obs_history）
+            runner.actor.class_name = (
+                "motrix_rl.rslrl.torch.models.cenet_actor:CENetActorModel")
+            runner.actor.hidden_dims = [512, 256, 128]
+            runner.actor.init_noise_std = 1.0
+
+            # Critic：标准 MLPModel，输入 privileged_obs
+            runner.critic.class_name = "MLPModel"
+            runner.critic.hidden_dims = [512, 256, 128]
+
+            # 观测分组：actor 用 policy+history，critic 用 privileged_obs
+            runner.obs_groups = {
+                "actor": ["policy", "obs_history"],
+                "critic": ["privileged_obs"],
+            }
+
     @rlcfg("go1-rough-terrain-walk")
     @dataclass
     class Go1WalkRoughRslrlPpo(Go1WalkFlatRslrlPpo):
@@ -123,3 +183,18 @@ class rslrl:
         def __post_init__(self):
             super().__post_init__()
             self.runner.experiment_name = "go1_stairs_terrain_walk"
+
+    @rlcfg("go1-stairs-terrain-walk-no-linevel")
+    @dataclass
+    class Go1WalkStairsNoLinvelRslrlPpo(Go1WalkRoughRslrlPpo):
+        """Go1 stairs terrain walk (no linear velocity obs) - RSLRL PPO config.
+
+        Uses [512, 256, 128] network from rough terrain config.
+        Increased iterations for 4-phase curriculum.
+        """
+
+        def __post_init__(self):
+            super().__post_init__()
+            self.runner.experiment_name = "go1_stairs_terrain_walk_no_linevel"
+            self.runner.max_iterations = 2000
+            self.runner.experiment_name = "go1_stairs_terrain_walk_no_linevel"