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

v1.0.5; add num_processes and support dual moe

Browse Source
This commit is contained in:
wty-yy
2026-01-07 00:20:22 +08:00
parent 9c3184a43d
commit 0f2dc29dfc
3 changed files with 24 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
robogauge/tasks/robots/go2/go2_moe.py
View File

@@ -23,13 +23,23 @@ class Go2MoE(Go2):
def get_action(self, obs: np.ndarray):
obs_tensor = torch.tensor(obs, dtype=torch.float32).unsqueeze(0).to(self.device)
action, results = self.model(obs_tensor)
if isinstance(results, tuple) and len(results) == 2:
weights, latent = results
if isinstance(results, tuple):
latent = results[-1]
latent = latent.detach().cpu().numpy().squeeze(0) if latent is not None else None
weights = weights.detach().cpu().numpy().squeeze(0) if weights is not None else None
student_weights = None
actor_weights = None
if len(results) >= 2:
student_weights = results[0]
student_weights = student_weights.detach().cpu().numpy().squeeze(0) if student_weights is not None else None
if len(results) >= 3:
actor_weights = results[1]
actor_weights = actor_weights.detach().cpu().numpy().squeeze(0) if actor_weights is not None else None
if self.cfg.control.save_additional_output:
self.save_info['latent'].append(latent)
self.save_info['weights'].append(weights)
if student_weights is not None:
self.save_info['weights'].append(student_weights)
if actor_weights is not None:
self.save_info['actor_weights'].append(actor_weights)
action = action.detach().cpu().numpy().squeeze(0)[self.model2mj_idx]
self.last_action = action
target_dof_pos = action * self.action_scale + self.default_dof_pos