From 5cbc8e2ca471ba975ce57340cccbf3ede234757f Mon Sep 17 00:00:00 2001 From: 8x54zj-m <8x54zj-m@motrixlab.local> Date: Wed, 1 Jul 2026 15:44:30 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20AdaBoot=20=E2=80=94=20adaptive=20veloci?= =?UTF-8?q?ty=20bootstrapping=20(paper=20Section=20II-C)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../rslrl/torch/models/cenet_actor.py | 36 +++++++++++++++++-- motrix_rl/src/motrix_rl/tasks/go1.py | 2 +- 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py b/motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py index cfb7317..7357e88 100644 --- a/motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py +++ b/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) diff --git a/motrix_rl/src/motrix_rl/tasks/go1.py b/motrix_rl/src/motrix_rl/tasks/go1.py index 681f5bd..fd5be29 100644 --- a/motrix_rl/src/motrix_rl/tasks/go1.py +++ b/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"], }