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

Update expertor for moe-cts, cts

Browse Source
This commit is contained in:
wty-yy
2026-01-02 23:05:37 +08:00
parent 102c0f8df7
commit c2519bb141
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
legged_gym/utils/exporter.py
View File

@@ -124,7 +124,7 @@ class _TorchPolicyExporter(torch.nn.Module):
self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1) self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
latent = self.student_encoder(self.history.flatten(1)) latent = self.student_encoder(self.history.flatten(1))
x = torch.cat([latent, x], dim=1) x = torch.cat([latent, x], dim=1)
return self.actor(x) return self.actor(x), (None, latent)
def forward_moe_cts(self, x): # x is single observations def forward_moe_cts(self, x): # x is single observations
x = self.normalizer(x) x = self.normalizer(x)
@@ -142,7 +142,7 @@ class _TorchPolicyExporter(torch.nn.Module):
x = torch.cat([latent, x], dim=1) x = torch.cat([latent, x], dim=1)
x_no_goal = torch.cat([latent, x_no_goal], dim=1) x_no_goal = torch.cat([latent, x_no_goal], dim=1)
mean_action, _, weights = self.actor(x, x_no_goal) mean_action, _, weights = self.actor(x, x_no_goal)
return mean_action, weights return mean_action, (weights, latent)
@torch.jit.export @torch.jit.export
def reset(self): def reset(self):