v0.1.3

robogauge/tasks/robots/go2/go2.py

@@ -0,0 +1,57 @@
+# -*- coding: utf-8 -*-
+'''
+@File    : go2.py
+@Time    : 2025/11/28 15:27:07
+@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.simulator.sim_data import SimData
+from robogauge.tasks.gauge.goal_data import GoalData
+from robogauge.utils.logger import logger
+
+class Go2(BaseRobot):
+    def __init__(self, cfg: Go2Config):
+        super().__init__(cfg)
+        self.max_velocity_cmd = np.array(cfg.control.max_velocity_cmd, dtype=np.float32)
+        self.default_dof_pos = np.array(cfg.control.default_dof_pos, dtype=np.float32)
+        self.last_action = np.zeros(self.num_action, dtype=np.float32)
+        self.action_scale = cfg.control.scales.action
+        self.mj2model_idx = self.cfg.control.mj2model_dof_indices
+        self.model2mj_idx = [self.mj2model_idx.index(i) for i in range(len(self.mj2model_idx))]
+    
+    def build_observation(self, sim_data: SimData, goal_data: GoalData) -> np.ndarray:
+        sim_proprio = sim_data.proprio
+        obs = np.zeros(self.num_obs)
+        if goal_data.goal_type == 'velocity':
+            ang_vel = sim_proprio.imu.ang_vel * self.cfg.control.scales.ang_vel
+            projected_gravity = get_projected_gravity(sim_proprio.imu.quat)
+            dof_pos = (sim_proprio.joint.pos - self.default_dof_pos) * self.cfg.control.scales.dof_pos
+            dof_vel = sim_proprio.joint.vel * self.cfg.control.scales.dof_vel
+
+            cmd = np.array(goal_data.velocity_goal.lin_vel[:2] + goal_data.velocity_goal.ang_vel[2:3], np.float32)
+            cmd = np.minimum(np.maximum(cmd, -self.max_velocity_cmd), self.max_velocity_cmd)
+            cmd *= self.cfg.control.scales.cmd
+
+            obs[:3] = ang_vel
+            obs[3:6] = projected_gravity
+            obs[6:9] = cmd
+            obs[9:9+self.num_action] = dof_pos[self.mj2model_idx]
+            obs[9+self.num_action:9+2*self.num_action] = dof_vel[self.mj2model_idx]
+            obs[9+2*self.num_action:9+3*self.num_action] = self.last_action[self.mj2model_idx]
+        else:
+            raise NotImplementedError(f"Goal type '{goal_data.goal_type}' not implemented in Go2 robot.")
+        return obs
+    
+    def get_action(self, obs: np.ndarray):
+        obs_tensor = torch.tensor(obs, dtype=torch.float32).unsqueeze(0).to(self.device)
+        action = self.model(obs_tensor).detach().cpu().numpy().squeeze(0)[self.model2mj_idx]
+        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_config.py

@@ -7,14 +7,44 @@
 @Blog    : https://wty-yy.github.io/
 @Desc    : Go2 Robot Configuration
 '''
+from typing_extensions import Literal
 from robogauge.tasks.robots import RobotConfig
 
 class Go2Config(RobotConfig):
+    robot_class = 'Go2'
+
     class assets:
         robot_xml = "{ROBOGAUGE_ROOT_DIR}/resources/robots/go2/go2.xml"
         robot_spawn_height = 0.1  # z [m]
 
-    class control:
+    class control(RobotConfig.control):
+        device = 'cpu'
+        torch_script_model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_cts_83501.pt"
         control_dt = 0.02  # 50 Hz
-        action_scale = 0.25  # scale for normalized actions
+        control_type = 'P'  # Position control
 
+        # Mujoco joint PD gains
+        p_gains = [20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0]  # [N*m/rad]
+        d_gains = [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]  # [N*m*s/rad]
+    
+        num_observations = 45
+        num_actions = 12
+
+        max_velocity_cmd = [1.5, 1.0, 2.0]
+        default_dof_pos = [0.1,  0.8,  -1.5,  -0.1, 0.8, -1.5, 
+                           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]
+
+        class scales(RobotConfig.control.scales):
+            lin_vel = 2.0
+            ang_vel = 0.25
+            dof_pos = 1.0
+            dof_vel = 0.05
+            action = 0.25  # target pos = action_scale * action * default_pos
+            cmd = [2.0, 2.0, 0.25]
+
+    class commands(RobotConfig.commands):
+        lin_vel_x = [-1, 1]  # min max [m/s]
+        lin_vel_y = [-1, 1]  # min max [m/s]
+        ang_vel_yaw = [-1, 1]  # min max [rad/s]