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

v0.1.5 prev1; Add ACMoE

Browse Source
This commit is contained in:
wty-yy
2026-01-06 23:38:32 +08:00
parent ddd7119050
commit 05e1e81d64
13 changed files with 601 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
legged_gym/utils/exporter.py
View File

@@ -85,6 +85,9 @@ class _TorchPolicyExporter(torch.nn.Module):
self.actor = copy.deepcopy(policy.actor_mcp)
self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
self.forward = self.forward_mcp_cts
elif hasattr(policy, "actor_moe"):
self.actor = copy.deepcopy(policy.actor_moe)
self.forward = self.forward_ac_moe
elif hasattr(policy, "actor"):
self.actor = copy.deepcopy(policy.actor)
if self.is_recurrent:
@@ -144,6 +147,14 @@ class _TorchPolicyExporter(torch.nn.Module):
mean_action, _, weights = self.actor(x, x_no_goal)
return mean_action, (weights, latent)
def forward_ac_moe(self, x): # x is single observations
x = self.normalizer(x)
self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
latent = self.student_encoder(self.history.flatten(1))
x = torch.cat([latent, x], dim=1)
mean, weights = self.actor(x)
return mean, (weights, latent)
@torch.jit.export
def reset(self):
if hasattr(self, 'history'):