v1.0.4; Add save_additional_output

robogauge/tasks/robots/go2/go2_config.py

@@ -7,7 +7,7 @@
 @Blog    : https://wty-yy.github.io/
 @Desc    : Go2 Robot Configuration
 '''
-from typing_extensions import Literal
+from typing import Literal
 from robogauge.tasks.robots import RobotConfig
 
 class Go2Config(RobotConfig):
@@ -37,6 +37,7 @@ class Go2Config(RobotConfig):
                            0.1,  1.0,  -1.5,  -0.1, 1.0, -1.5] 
         
         mj2model_dof_indices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
+        save_additional_output = False
 
         class scales(RobotConfig.control.scales):
             lin_vel = 2.0

robogauge/tasks/robots/go2/go2_moe.py

@@ -9,20 +9,40 @@
 '''
 import torch
 import numpy as np
+from collections import defaultdict
 
 from robogauge.tasks.robots.go2.go2 import Go2
+from robogauge.utils.logger import logger
 
 class Go2MoE(Go2):
+    def __init__(self, cfg):
+        super().__init__(cfg)
+        self.save_info = defaultdict(list)
+        self.save_count = 0
+
     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):
+        if isinstance(results, tuple) and len(results) == 2:
             weights, latent = results
-            latent = latent.detach().cpu().numpy().squeeze(0)
-        else:
-            weights = results
+            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
+            if self.cfg.control.save_additional_output:
+                self.save_info['latent'].append(latent)
+                self.save_info['weights'].append(weights)
         action = action.detach().cpu().numpy().squeeze(0)[self.model2mj_idx]
-        weights = weights.detach().cpu().numpy().squeeze(0)
         self.last_action = action
         target_dof_pos = action * self.action_scale + self.default_dof_pos
         return target_dof_pos, self.p_gains, self.d_gains, self.control_type
+
+    def reset(self):
+        super().reset()
+        save_path = logger.log_dir / f"moe_info_{self.save_count}.npz"
+        if self.cfg.control.save_additional_output:
+            np.savez_compressed(save_path,
+                weights=np.array(self.save_info['weights']),
+                latent=np.array(self.save_info['latent'])
+            )
+            logger.info(f"Saved MoE info to {save_path}")
+        self.save_count += 1
+        self.save_info = defaultdict(list)

robogauge/tasks/robots/go2/go2_moe_config.py

@@ -14,6 +14,7 @@ class Go2MoEConfig(Go2Config):
 
     class control(Go2Config.control):
         model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_moe_cts_124k.pt"
+        save_additional_output = False
 
 class Go2MoETerrainConfig(Go2MoEConfig):
     """ Go2 MoE Robot Configuration for Terrain Tasks (wave, stairs up/down, slope, obstacles) """