v0.1.7

robogauge/tasks/robots/__init__.py

@@ -2,3 +2,5 @@ from .base_robot_config import RobotConfig
 from .base_robot import BaseRobot
 from .go2.go2_config import Go2Config
 from .go2.go2 import Go2
+from .go2.go2_moe_config import Go2MoEConfig
+from .go2.go2_moe import Go2MoE

robogauge/tasks/robots/go2/go2_moe.py

@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+'''
+@File    : go2_moe.py
+@Time    : 2025/12/05 17:29:16
+@Author  : wty-yy
+@Version : 1.0
+@Blog    : https://wty-yy.github.io/
+@Desc    : None
+'''
+import torch
+import numpy as np
+
+from robogauge.tasks.robots.base_robot import BaseRobot, get_projected_gravity
+from robogauge.tasks.robots.go2.go2_config import Go2Config
+from robogauge.tasks.robots.go2.go2 import Go2
+from robogauge.tasks.simulator.sim_data import SimData
+from robogauge.tasks.gauge.goal_data import GoalData
+from robogauge.utils.logger import logger
+
+class Go2MoE(Go2):
+    def get_action(self, obs: np.ndarray):
+        obs_tensor = torch.tensor(obs, dtype=torch.float32).unsqueeze(0).to(self.device)
+        action, weights = self.model(obs_tensor)
+        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

robogauge/tasks/robots/go2/go2_moe_config.py

@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+'''
+@File    : go2_moe_config.py
+@Time    : 2025/12/05 17:32:27
+@Author  : wty-yy
+@Version : 1.0
+@Blog    : https://wty-yy.github.io/
+@Desc    : None
+'''
+from robogauge.tasks.robots.go2.go2_config import Go2Config
+
+class Go2MoEConfig(Go2Config):
+    robot_class = 'Go2MoE'
+
+    class control(Go2Config.control):
+        model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_moe_cts_124k.pt"