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

v0.1.8; add configs, fix bugs

Browse Source
This commit is contained in:
wty-yy
2026-01-10 02:50:58 +08:00
parent 836a4e544f
commit bcb3757403
7 changed files with 711 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
legged_gym/utils/exporter.py
View File

@@ -210,10 +210,13 @@ class _OnnxPolicyExporter(torch.nn.Module):
elif hasattr(policy, "student_moe_encoder"):
self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder)
self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
self.history_length = policy.history.shape[1]
self.forward = self.forward_moe_cts
self.input_dim = self.history_length * policy.history.shape[2]
if hasattr(policy, "obs_no_goal_mask"):
self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
else:
self.forward = self.forward_rem_cts
else: # PPO
self.forward = self.forward_ppo
@@ -286,6 +289,17 @@ class _OnnxPolicyExporter(torch.nn.Module):
x = torch.cat([latent, last_obs], dim=1)
return self.actor(x), weights, latent
def forward_rem_cts(self, x):
x = self.normalizer(x)
history, obs_dim = self.flatten_obs(x)
last_obs = history[:, -obs_dim:]
latent, weights = self.student_moe_encoder(history)
x = torch.cat([latent, last_obs], dim=1)
return self.actor(x), weights, latent
def forward_mcp_cts(self, x):
x = self.normalizer(x)