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

11
scripts/export_dreamwaq_rsl_onnx.py
View File

@@ -22,13 +22,20 @@ PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
class DwaqInfer(tnn.Module):
"""Deterministic inference: CENet mean code + actor."""
"""Deterministic inference: CENet mean code + RunningStats normalize + actor."""
def __init__(self, ac: ActorCritic_DWAQ):
super().__init__()
self.encoder = ac.encoder
self.encode_mean_vel = ac.encode_mean_vel
self.encode_mean_latent = ac.encode_mean_latent
self.actor = ac.actor
# Bake RunningStats into the exported model
if hasattr(ac, 'actor_normalizer') and ac.actor_normalizer.count > 10:
self.register_buffer("norm_mean", ac.actor_normalizer.mean.clone())
self.register_buffer("norm_std", ac.actor_normalizer.var.sqrt().clone())
self.has_norm = True
else:
self.has_norm = False
def forward(self, obs, obs_history):
h = self.encoder(obs_history.reshape(obs_history.shape[0], -1)) # (B,225)->(B,64)
@@ -36,6 +43,8 @@ class DwaqInfer(tnn.Module):
latent = self.encode_mean_latent(h) # (B,16) mean latent
code = torch.cat([vel, latent], dim=-1) # (B,19) = [vel, latent]
x = torch.cat([code, obs], dim=-1) # (B,64) = [code, obs]
if self.has_norm:
x = torch.clamp((x - self.norm_mean) / (self.norm_std + 1e-8), -5.0, 5.0)
return self.actor(x) # (B,12)