cyy/Motrixlab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Backup DreamWaQ rslrl stability fixes

Browse Source
This commit is contained in:
8x54zj-m
2026-07-22 02:17:01 +08:00
parent c664e7422f
commit 3648551043
17 changed files with 1068 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
motrix_rl/src/motrix_rl/dwaq_rsl/actor_critic_dwaq.py
View File

@@ -4,6 +4,43 @@ import torch
import torch.nn as nn
from torch.distributions import Normal
class RunningStats(nn.Module):
"""Running mean/variance tracker for actor observation normalization.
Matches SKRL RunningStandardScaler behavior. Without it, actor sees
unscaled inputs at inference → garbage actions → robot collapses.
"""
def __init__(self, num_features: int, eps: float = 1e-8, clip: float = 5.0):
super().__init__()
self.eps = eps
self.clip = clip
self.register_buffer("count", torch.zeros(1, dtype=torch.int64))
self.register_buffer("mean", torch.zeros(num_features))
self.register_buffer("var", torch.ones(num_features))
@torch.no_grad()
def update(self, x):
if x.dim() != 2:
return
n = x.shape[0]
batch_mean = x.mean(dim=0)
batch_var = x.var(dim=0, unbiased=False)
batch_count = torch.tensor(n, dtype=torch.int64, device=x.device)
delta = batch_mean - self.mean
total = self.count + batch_count
self.mean.add_(delta * batch_count.float() / total.float())
m_a = self.var * self.count.float()
m_b = batch_var * batch_count.float()
m2 = m_a + m_b + delta.pow(2) * self.count.float() * batch_count.float() / total.float()
self.var.copy_(m2 / total.float())
self.count.add_(batch_count)
def forward(self, x):
return torch.clamp(
(x - self.mean) / (self.var.sqrt() + self.eps), -self.clip, self.clip)
class ActorCritic_DWAQ(nn.Module):
def __init__(self, num_actor_obs, num_critic_obs, num_actions, cenet_in_dim, cenet_out_dim, activation="elu", init_noise_std=1.0,):
super().__init__()
@@ -52,6 +89,7 @@ class ActorCritic_DWAQ(nn.Module):
)
self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
self.actor_normalizer = RunningStats(actor_input_dim, eps=1e-8, clip=5.0)
self.distribution = None
# disable args validation for speedup
Normal.set_default_validate_args = False
@@ -89,7 +127,7 @@ class ActorCritic_DWAQ(nn.Module):
# code = mean_latent + var*code_temp
# print("latent : ",code[0])
mean_vel = self.encode_mean_vel(distribution)
logvar_vel = self.encode_mean_vel(distribution)
logvar_vel = self.encode_logvar_vel(distribution) # FIXED: was encode_mean_vel
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)
@@ -117,17 +155,23 @@ class ActorCritic_DWAQ(nn.Module):
self.distribution = Normal(mean, mean * 0.0 + self.std)
def act(self, observations, obs_history, **kwargs):
code,_,decode,_,_,_,_ = self.cenet_forward(obs_history)
observations = torch.cat((code,observations),dim=-1)
code, _, decode, _, _, _, _ = self.cenet_forward(obs_history)
observations = torch.cat((code, observations), dim=-1)
if self.training:
self.actor_normalizer.update(observations)
if self.actor_normalizer.count > 10:
observations = self.actor_normalizer(observations)
self.update_distribution(observations)
return self.distribution.sample()
def get_actions_log_prob(self, actions):
return self.distribution.log_prob(actions).sum(dim=-1)
def act_inference(self, observations,obs_history):
code,_,decode,_,_,_,_ = self.cenet_forward(obs_history)
observations = torch.cat((code,observations),dim=-1)
def act_inference(self, observations, obs_history):
code, _, decode, _, _, _, _ = self.cenet_forward(obs_history)
observations = torch.cat((code, observations), dim=-1)
if self.actor_normalizer.count > 10:
observations = self.actor_normalizer(observations)
actions_mean = self.actor(observations)
return actions_mean

2
motrix_rl/src/motrix_rl/dwaq_rsl/env_wrapper.py
View File

@@ -20,7 +20,7 @@ import torch
class DwaqVecEnv:
"""Wraps a MotrixLab DreamWaQ NpEnv for the upstream DreamWaQ rsl_rl runner."""
def __init__(self, env, device, num_obs=45, num_privileged_obs=235,
def __init__(self, env, device, num_obs=45, num_privileged_obs=247,
num_obs_hist=5, num_actions=12, clip_actions=100.0):
self._env = env
self.device = device

17
motrix_rl/src/motrix_rl/dwaq_rsl/runner.py
View File

@@ -158,12 +158,19 @@ class OnPolicyRunner:
f"rew={rew:.2f} eplen={elen:.0f}")
def save(self, path, infos=None):
torch.save({
save_dict = {
"model_state_dict": self.alg.actor_critic.state_dict(),
"optimizer_state_dict": self.alg.optimizer.state_dict(),
"iter": self.current_learning_iteration,
"infos": infos,
}, path)
}
# Persist RunningStats normalizer state for inference
if hasattr(self.alg.actor_critic, 'actor_normalizer'):
n = self.alg.actor_critic.actor_normalizer
save_dict["normalizer_count"] = n.count
save_dict["normalizer_mean"] = n.mean
save_dict["normalizer_var"] = n.var
torch.save(save_dict, path)
def load(self, path, load_optimizer=True):
loaded_dict = torch.load(path, map_location=self.device)
@@ -171,6 +178,12 @@ class OnPolicyRunner:
if load_optimizer:
self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"])
self.current_learning_iteration = loaded_dict["iter"]
# Restore RunningStats normalizer state
if "normalizer_count" in loaded_dict and hasattr(self.alg.actor_critic, 'actor_normalizer'):
n = self.alg.actor_critic.actor_normalizer
n.count.copy_(loaded_dict["normalizer_count"])
n.mean.copy_(loaded_dict["normalizer_mean"])
n.var.copy_(loaded_dict["normalizer_var"])
return loaded_dict["infos"]
def get_inference_policy(self, device=None):

39
motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py
View File

@@ -14,6 +14,7 @@
import torch
import torch.nn as nn
from torch.distributions import Normal
from tensordict import TensorDict
from rsl_rl.models.mlp_model import MLPModel
@@ -113,23 +114,43 @@ class CENetActorModel(MLPModel):
# 在 nn.Module.__init__ 之后创建 VAE 子模块
self.vae = CENetVAE(cenet_in_dim, cenet_out_dim, activation)
self.action_clip = 4.0
self.std_clip = 0.6
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）。"""
# 先 clamp std，再调父类创建分布
"""覆盖父类 — 限制动作分布，避免采样动作进入 PD/关节限位饱和区。"""
mean = torch.clamp(self.mlp(obs), -self.action_clip, self.action_clip)
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)
self.std.clamp_(min=1e-6, max=self.std_clip)
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)
self.log_std.clamp_(min=-20.0, max=torch.log(torch.tensor(self.std_clip)).item())
if self.noise_std_type == "scalar":
std = self.std.expand_as(mean)
elif self.noise_std_type == "log":
std = torch.exp(self.log_std).expand_as(mean)
else:
raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}.")
else:
std = torch.full_like(mean, self.std_clip)
self.distribution = Normal(mean, std)
def forward(
self,
obs: TensorDict,
masks: torch.Tensor | None = None,
hidden_state: HiddenState = None,
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)
def _get_latent_dim(self) -> int:
"""Actor 实际输入：code(19) + policy(45) = 64。
@@ -164,10 +185,10 @@ class CENetActorModel(MLPModel):
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)
if torch.isnan(policy_obs).any():
policy_obs = torch.nan_to_num(policy_obs, nan=0.0)
if not torch.isfinite(code).all():
code = torch.nan_to_num(code, nan=0.0, posinf=0.0, neginf=0.0)
if not torch.isfinite(policy_obs).all():
policy_obs = torch.nan_to_num(policy_obs, nan=0.0, posinf=0.0, neginf=0.0)
latent = torch.cat([code, policy_obs], dim=-1) # (N, 64)
return latent

47
motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py
View File

@@ -81,8 +81,17 @@ class DreamWaQPPO(PPO):
returns_batch, old_actions_log_prob_batch, old_mu_batch,
old_sigma_batch, hid_states_batch, masks_batch,
) in generator:
# NaN 检测 — 数据有 NaN 就跳过这个 batch
if torch.isnan(obs_batch["policy"]).any() or torch.isnan(obs_batch["obs_history"]).any():
# 多层有限值检测 — 数据异常就跳过这个 batch，避免污染网络参数。
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 not torch.isfinite(actions_batch).all()
or not torch.isfinite(target_values_batch).all()
or not torch.isfinite(advantages_batch).all()
or not torch.isfinite(returns_batch).all()
or not torch.isfinite(old_actions_log_prob_batch).all()
or not torch.isfinite(old_mu_batch).all()
or not torch.isfinite(old_sigma_batch).all()):
continue
# ── 标准 PPO 前向（各调用一次，复用结果）──
self.actor(obs_batch, masks=masks_batch, stochastic_output=True)
@@ -106,6 +115,11 @@ class DreamWaQPPO(PPO):
actions_batch, actions_log_prob_batch,
old_actions_log_prob_batch, advantages_batch)
# NaN guard: 损失值异常 → 跳过此 batch，保护模型权重
if (not torch.isfinite(surrogate_loss)
or not torch.isfinite(value_loss)):
continue
# ── KL 自适应 schedule ──
if self.desired_kl is not None and self.schedule == 'adaptive':
with torch.inference_mode():
@@ -136,20 +150,41 @@ class DreamWaQPPO(PPO):
- self.entropy_coef * entropy_loss
+ autoenc_loss
)
if not torch.isfinite(loss):
continue
# ── 梯度更新 ──
self.optimizer.zero_grad()
loss.backward()
actor_grads_ok = all(
p.grad is None or torch.isfinite(p.grad).all()
for p in self.actor.parameters()
)
critic_grads_ok = all(
p.grad is None or torch.isfinite(p.grad).all()
for p in self.critic.parameters()
)
if not (actor_grads_ok and critic_grads_ok):
self.optimizer.zero_grad(set_to_none=True)
continue
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)
try:
nn.utils.clip_grad_norm_(
self.actor.parameters(), self.max_grad_norm, error_if_nonfinite=True)
nn.utils.clip_grad_norm_(
self.critic.parameters(), self.max_grad_norm, error_if_nonfinite=True)
except RuntimeError:
self.optimizer.zero_grad(set_to_none=True)
continue
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.nan_to_num_(nan=0.5, posinf=1.0, neginf=1.0)
self.actor.std.clamp_(min=1e-6)
if not torch.isfinite(self.actor.std).all():
self.actor.std.data.fill_(0.5)
std_clip = getattr(self.actor, "std_clip", 0.6)
self.actor.std.data.clamp_(min=1e-6, max=std_clip)
# ── 累计日志 ──
mean_value_loss += value_loss.item()

6
motrix_rl/src/motrix_rl/skrl/jax/train/dreamwaq_ppo.py
View File

@@ -4,7 +4,7 @@
CENet (VAE): history(225) → [128,64] → latent(16) + vel_est(3) = code(19)
Decoder: code(19) → [64,128] → next_obs(45)
Actor: code(19) + obs(45) = 64 → [512,256,128] → action(12)
Critic: privileged_obs(235) → [512,256,128] → value(1)
Critic: privileged_obs(247) → [512,256,128] → value(1)
VAE Loss: reconstruction_MSE + velocity_MSE + beta * KL
"""
@@ -189,7 +189,7 @@ class DreamWaQWrapper:
@property
def privileged_obs(self):
return self._env._state.info.get("privileged_obs",
np.zeros((self._num_envs, 235), dtype=np.float32))
np.zeros((self._num_envs, 247), dtype=np.float32))
@property
def obs_history(self):
@@ -398,7 +398,7 @@ class DreamWaQTrainer:
def __call__(self, inputs, role):
kernel_init = nn.initializers.orthogonal(jnp.sqrt(2))
x = inputs["states"]
# Critic: obs(45) + base_vel(3) + heights(187) = 235
# Critic: obs(45) + base_vel(3) + heights(187) = 247
# Layout: [code(19) | obs(45) | base_vel(3) | heights(187)]
x_c = jnp.concatenate([x[:, 19:64], x[:, 64:254]], axis=-1)
for d in value_cfg.hiddens:

2
motrix_rl/src/motrix_rl/tasks/go1.py
View File

@@ -112,7 +112,7 @@ class rslrl:
# Runner 设置（严格对齐上游 LeggedRobotCfgPPO + Go1RoughCfgPPO）
runner.seed = 5 # 上游 seed=5
runner.max_iterations = 3000 # 续训到 3000 轮
runner.max_iterations = 5000
runner.num_steps_per_env = 24
runner.experiment_name = "go1_dreamwaq_walk"
runner.save_interval = 50