feat: AdaBoot — adaptive velocity bootstrapping (paper Section II-C)

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

@@ -114,6 +114,9 @@ class CENetActorModel(MLPModel):
         # 在 nn.Module.__init__ 之后创建 VAE 子模块
         self.vae = CENetVAE(cenet_in_dim, cenet_out_dim, activation)
         self._last_cenet_output = None
+        # AdaBoot: 自适应 bootstrapping（论文 Section II-C）
+        self._adaboot_cv_buffer = []  # 速度估计误差的 CV 历史
+        self._adaboot_prob = 1.0      # 当前 bootstrap 概率（1.0 = 完全信任 GT）
 
     def _update_distribution(self, obs: torch.Tensor) -> None:
         """覆盖父类 — 强制 std > 0 再创建 Normal 分布（防止 NaN）。"""
@@ -129,12 +132,22 @@ class CENetActorModel(MLPModel):
         super()._update_distribution(obs)
 
     def _get_latent_dim(self) -> int:
-        """Actor 实际输入：code(19) + policy(45) = 64。"""
-        return self.obs_dim - self._history_dim + self._code_dim
+        """Actor 实际输入：code(19) + policy(45) = 64。
+
+        privileged_obs 只用于 AdaBoot，不计入 latent dim。
+        """
+        # obs_groups 可能包含 privileged_obs（AdaBoot），需排除
+        extra_dims = 0
+        if "privileged_obs" in self.obs_groups:
+            extra_dims = 247  # privileged_obs 维度
+        return self.obs_dim - self._history_dim - extra_dims + 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)。"""
+        """提取观测 → VAE 编码 → 拼接 code + policy → 返回 latent(64)。
+
+        AdaBoot: 训练时以概率 p 用 ground truth velocity 替换 CENet 的速度估计。
+        """
         policy_obs = obs["policy"]           # (N, 45)
         obs_history = obs["obs_history"]     # (N, 225)
 
@@ -142,6 +155,23 @@ class CENetActorModel(MLPModel):
         self._last_cenet_output = out
         code, code_vel, decode, mean_vel, logvar_vel, mean_latent, logvar_latent = out
 
+        # ── AdaBoot: 自适应速度 Bootstrap（论文 Section II-C）──
+        if self.training and "privileged_obs" in obs:
+            gt_vel = obs["privileged_obs"][:, 45:48]  # ground truth body velocity
+            # 先计算 CENet 原始误差的 CV（替换前）
+            with torch.no_grad():
+                raw_error = code_vel - gt_vel
+                cv = raw_error.std() / (raw_error.abs().mean() + 1e-6)
+                self._adaboot_cv_buffer.append(cv.item())
+                if len(self._adaboot_cv_buffer) > 1000:
+                    self._adaboot_cv_buffer = self._adaboot_cv_buffer[-1000:]
+                mean_cv = sum(self._adaboot_cv_buffer) / len(self._adaboot_cv_buffer)
+                self._adaboot_prob = max(0.0, min(1.0, mean_cv * 5.0))
+            # 以概率 _adaboot_prob 用 GT 替换 CENet 的速度估计
+            mask = torch.rand(code_vel.shape[0], 1, device=code_vel.device) < self._adaboot_prob
+            code_vel = torch.where(mask, gt_vel, code_vel)
+            code = torch.cat([code_vel, code[:, 3:]], dim=-1)
+
         # 防止 VAE NaN 传播到下游
         if torch.isnan(code).any():
             code = torch.nan_to_num(code, nan=0.0)

motrix_rl/src/motrix_rl/tasks/go1.py

@@ -144,7 +144,7 @@ class rslrl:
 
             # 观测分组：actor 用 policy+history，critic 用 privileged_obs
             runner.obs_groups = {
-                "actor": ["policy", "obs_history"],
+                "actor": ["policy", "obs_history", "privileged_obs"],  # privileged_obs 用于 AdaBoot
                 "critic": ["privileged_obs"],
             }