v0.1.3

README.md

@@ -44,3 +44,21 @@
 | 4 | base高度变化 | 固定高度 | 高速移动时机身存在趴低问题 | 平地 |
 | 5 | 速度对角突变base高度变化 | 固定高度 | 速度发生对角突变时无法平衡 | 平地 |
 | 6 | 高速移动急停稳定性 | 固定用时 | 楼梯上静止时, 关节不稳定 | Any |
+
+## 创建新任务
+评测任务注册在[`tasks/__init__.py`](./robogauge/tasks/__init__.py)中完成, 包含四个部分:
+- BasePipline: 标准Pipline一般无需修改
+- MujocoConfig: 仿真器配置文件, 一般无需修改
+- BaseGaugeConfig: 指标配置文件, 参考下文创建新指标
+- RobotConfig: 机器人配置文件, 参考下文创建新机器人
+### 新指标
+在`robogauge/tasks/gauge`下创建新的场景、评估指标
+### 新机器人
+在`robogauge/tasks/robots`下创建新机器人、配置控制模型, 参考`go2`配置
+- [go2.py](./robogauge/tasks/robots/go2/go2.py)控制模型, 包含观测构建和动作输出两个函数, 继承`BaseRobot`
+- [go2_config.py](./robogauge/tasks/robots/go2/go2_config.py)配置文件, 包含机器人模型xml路径, 控制模型路径, 观测构建参数, 关节映射顺序, 各类缩放系数, 控制频率等
+
+## 注意事项
+### 导入新机器人/控制模型
+1. 在Robot中创建新的机器人xml文件时, 需包含力矩控制`actuator`, 传感器`sensor - jointpos, jointvel, imu (framequat, gyro, accelerometer)`, 参考[`go2.xml`](resources/robots/go2/go2.xml), **注意: actuator的顺序需要和joint顺序一致, 该顺序称为mujoco关节顺序**
+2. 控制模型相关位置位于[`RobotConfig.control`](robogauge/tasks/robots/base_robot_config.py)中, 不同仿真中关节顺序可能不同, 需保证`mj2model_dof_indices`从mujoco映射到模型训练的关节次序配置正确 (IsaacGym次序和Mujoco相同), 其他模型配置需保持一致

UPDATE.md

@@ -1,4 +1,7 @@
 # UPDATE
+## 20251128
+### v0.1.3
+1. 完成go2模型预测
 
 ## 20251127
 ### v0.1.2

resources/robots/go2/go2.xml

@@ -40,7 +40,6 @@
     <material name="metal" rgba=".9 .95 .95 1" />
     <material name="black" rgba="0 0 0 1" />
     <material name="white" rgba="1 1 1 1" />
-    <!-- <material name="white" rgba="1 1 0 1" /> -->
     <material name="gray" rgba="0.671705 0.692426 0.774270 1" />
 
     <mesh file="base_0.obj" />
@@ -222,59 +221,60 @@
   </worldbody>
 
   <actuator>
-    <motor class="abduction" name="FR_hip" joint="FR_hip_joint" />
+    <!-- mujoco dof actuator (not real robot actuators order) -->
-    <motor class="hip" name="FR_thigh" joint="FR_thigh_joint" />
+    <motor name="FL_hip"   joint="FL_hip_joint"/>
-    <motor class="knee" name="FR_calf" joint="FR_calf_joint" />
+    <motor name="FL_thigh" joint="FL_thigh_joint"/>
-    <motor class="abduction" name="FL_hip" joint="FL_hip_joint" />
+    <motor name="FL_calf"  joint="FL_calf_joint"/>
-    <motor class="hip" name="FL_thigh" joint="FL_thigh_joint" />
+    <motor name="FR_hip"   joint="FR_hip_joint"/>
-    <motor class="knee" name="FL_calf" joint="FL_calf_joint" />
+    <motor name="FR_thigh" joint="FR_thigh_joint"/>
-    <motor class="abduction" name="RR_hip" joint="RR_hip_joint" />
+    <motor name="FR_calf"  joint="FR_calf_joint"/>
-    <motor class="hip" name="RR_thigh" joint="RR_thigh_joint" />
+    <motor name="RL_hip"   joint="RL_hip_joint"/>
-    <motor class="knee" name="RR_calf" joint="RR_calf_joint" />
+    <motor name="RL_thigh" joint="RL_thigh_joint"/>
-    <motor class="abduction" name="RL_hip" joint="RL_hip_joint" />
+    <motor name="RL_calf"  joint="RL_calf_joint"/>
-    <motor class="hip" name="RL_thigh" joint="RL_thigh_joint" />
+    <motor name="RR_hip"   joint="RR_hip_joint"/>
-    <motor class="knee" name="RL_calf" joint="RL_calf_joint" />
+    <motor name="RR_thigh" joint="RR_thigh_joint"/>
+    <motor name="RR_calf"  joint="RR_calf_joint"/>
   </actuator>
 
   <sensor>
-    <jointpos name="FR_hip_pos" joint="FR_hip_joint" />
-    <jointpos name="FR_thigh_pos" joint="FR_thigh_joint" />
-    <jointpos name="FR_calf_pos" joint="FR_calf_joint" />
     <jointpos name="FL_hip_pos" joint="FL_hip_joint" />
     <jointpos name="FL_thigh_pos" joint="FL_thigh_joint" />
     <jointpos name="FL_calf_pos" joint="FL_calf_joint" />
-    <jointpos name="RR_hip_pos" joint="RR_hip_joint" />
+    <jointpos name="FR_hip_pos" joint="FR_hip_joint" />
-    <jointpos name="RR_thigh_pos" joint="RR_thigh_joint" />
+    <jointpos name="FR_thigh_pos" joint="FR_thigh_joint" />
-    <jointpos name="RR_calf_pos" joint="RR_calf_joint" />
+    <jointpos name="FR_calf_pos" joint="FR_calf_joint" />
     <jointpos name="RL_hip_pos" joint="RL_hip_joint" />
     <jointpos name="RL_thigh_pos" joint="RL_thigh_joint" />
     <jointpos name="RL_calf_pos" joint="RL_calf_joint" />
+    <jointpos name="RR_hip_pos" joint="RR_hip_joint" />
+    <jointpos name="RR_thigh_pos" joint="RR_thigh_joint" />
+    <jointpos name="RR_calf_pos" joint="RR_calf_joint" />
 
-    <jointvel name="FR_hip_vel" joint="FR_hip_joint" />
-    <jointvel name="FR_thigh_vel" joint="FR_thigh_joint" />
-    <jointvel name="FR_calf_vel" joint="FR_calf_joint" />
     <jointvel name="FL_hip_vel" joint="FL_hip_joint" />
     <jointvel name="FL_thigh_vel" joint="FL_thigh_joint" />
     <jointvel name="FL_calf_vel" joint="FL_calf_joint" />
-    <jointvel name="RR_hip_vel" joint="RR_hip_joint" />
+    <jointvel name="FR_hip_vel" joint="FR_hip_joint" />
-    <jointvel name="RR_thigh_vel" joint="RR_thigh_joint" />
+    <jointvel name="FR_thigh_vel" joint="FR_thigh_joint" />
-    <jointvel name="RR_calf_vel" joint="RR_calf_joint" />
+    <jointvel name="FR_calf_vel" joint="FR_calf_joint" />
     <jointvel name="RL_hip_vel" joint="RL_hip_joint" />
     <jointvel name="RL_thigh_vel" joint="RL_thigh_joint" />
     <jointvel name="RL_calf_vel" joint="RL_calf_joint" />
+    <jointvel name="RR_hip_vel" joint="RR_hip_joint" />
+    <jointvel name="RR_thigh_vel" joint="RR_thigh_joint" />
+    <jointvel name="RR_calf_vel" joint="RR_calf_joint" />
 
-    <jointactuatorfrc name="FR_hip_torque" joint="FR_hip_joint" noise="0.01" />
-    <jointactuatorfrc name="FR_thigh_torque" joint="FR_thigh_joint" noise="0.01" />
-    <jointactuatorfrc name="FR_calf_torque" joint="FR_calf_joint" noise="0.01" />
     <jointactuatorfrc name="FL_hip_torque" joint="FL_hip_joint" noise="0.01" />
     <jointactuatorfrc name="FL_thigh_torque" joint="FL_thigh_joint" noise="0.01" />
     <jointactuatorfrc name="FL_calf_torque" joint="FL_calf_joint" noise="0.01" />
-    <jointactuatorfrc name="RR_hip_torque" joint="RR_hip_joint" noise="0.01" />
+    <jointactuatorfrc name="FR_hip_torque" joint="FR_hip_joint" noise="0.01" />
-    <jointactuatorfrc name="RR_thigh_torque" joint="RR_thigh_joint" noise="0.01" />
+    <jointactuatorfrc name="FR_thigh_torque" joint="FR_thigh_joint" noise="0.01" />
-    <jointactuatorfrc name="RR_calf_torque" joint="RR_calf_joint" noise="0.01" />
+    <jointactuatorfrc name="FR_calf_torque" joint="FR_calf_joint" noise="0.01" />
     <jointactuatorfrc name="RL_hip_torque" joint="RL_hip_joint" noise="0.01" />
     <jointactuatorfrc name="RL_thigh_torque" joint="RL_thigh_joint" noise="0.01" />
     <jointactuatorfrc name="RL_calf_torque" joint="RL_calf_joint" noise="0.01" />
+    <jointactuatorfrc name="RR_hip_torque" joint="RR_hip_joint" noise="0.01" />
+    <jointactuatorfrc name="RR_thigh_torque" joint="RR_thigh_joint" noise="0.01" />
+    <jointactuatorfrc name="RR_calf_torque" joint="RR_calf_joint" noise="0.01" />
 
     <framequat name="imu_quat" objtype="site" objname="imu" />
     <gyro name="imu_gyro" site="imu" />
@@ -283,9 +283,4 @@
     <framepos name="frame_pos" objtype="site" objname="imu" />
     <framelinvel name="frame_vel" objtype="site" objname="imu" />
   </sensor>
-
-  <!-- <keyframe>
-    <key name="home" qpos="0 0 0.27 1 0 0 0 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8"
-      ctrl="0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8 0 0.9 -1.8" />
-  </keyframe> -->
 </mujoco>

robogauge/__init__.py

@@ -1,4 +1,4 @@
 from pathlib import Path
 
 __version__ = "0.1.0"
-ROBOGAUGE_ROOT_DIR = str(Path(__file__).parent)
+ROBOGAUGE_ROOT_DIR = str(Path(__file__).parents[1])

robogauge/scripts/run.py

@@ -7,6 +7,9 @@
 @Blog    : https://wty-yy.github.io/
 @Desc    : Run Robogauge Pipeline
 '''
+import os
+os.environ['MUJOCO_GL'] = 'glfw'  # avoid mujoco.Renderer EGL context error
+
 from robogauge.tasks import *
 from robogauge.utils.task_register import task_register
 from robogauge.utils.helpers import parse_args

robogauge/tasks/__init__.py

@@ -1,7 +1,8 @@
 from robogauge.utils.task_register import task_register
 from robogauge.tasks.simulator.mujoco_config import MujocoConfig
-from robogauge.tasks.robots import RobotConfig
+from robogauge.tasks.robots import RobotConfig, Go2Config
 from robogauge.tasks.pipeline import BasePipeline
 from robogauge.tasks.gauge import BaseGaugeConfig
 
 task_register.register('base', BasePipeline, MujocoConfig, BaseGaugeConfig, RobotConfig)
+task_register.register('go2', BasePipeline, MujocoConfig, BaseGaugeConfig, Go2Config)

robogauge/tasks/gauge/base_gauge.py

@@ -9,6 +9,9 @@
 '''
 from robogauge.tasks.robots.base_robot_config import RobotConfig
 from robogauge.tasks.gauge.base_gauge_config import BaseGaugeConfig
+from robogauge.tasks.gauge.goal_data import GoalData, VelocityGoal, PositionGoal
+from robogauge.tasks.simulator.sim_data import SimData
+from robogauge.utils.logger import logger
 
 class BaseGauge:
     def __init__(self, cfg: BaseGaugeConfig):
@@ -20,10 +23,19 @@ class BaseGauge:
     def is_done(self) -> bool:
         return False
     
-    def get_goal(self) -> dict:
+    def get_goal(self) -> GoalData:
-        goal = {}
+        goal = GoalData(
+            goal_type='velocity',
+            velocity_goal=VelocityGoal(
+                lin_vel=[5.0, 0.0, 0.0],
+                ang_vel=[0.0, 0.0, 0.0]
+            )
+        )
         return goal
     
-    def update_metrics(self, sim_info: dict):
+    def update_metrics(self, sim_data: SimData):
-        ...
+        if sim_data.n_step % int(0.1 / sim_data.sim_dt) != 0:
+            return
+        for i in range(len(sim_data.proprio.joint.force)):
+            logger.log(sim_data.proprio.joint.force[i], f'dof/force_{i}', step=sim_data.n_step)
     

robogauge/tasks/gauge/goal_data.py

@@ -0,0 +1,22 @@
+from enum import Enum
+from dataclasses import dataclass
+from typing import List, Optional, Literal
+
+@dataclass
+class VelocityGoal:
+    lin_vel: List[float]  # x, y, z [m/s], z is ignored for ground robots
+    ang_vel: List[float]  # roll, pitch, yaw [rad/s], roll and pitch are ignored for ground robots
+
+@dataclass
+class PositionGoal:
+    # relative to robot's current position
+    target_pos: List[float]  # x, y, z [m], z is ignored for ground robots
+    # reach target orientation
+    target_quat: List[float]  # x, y, z, w quaternion
+    tolerance: float  # [m] position tolerance to consider goal reached
+
+@dataclass
+class GoalData:
+    goal_type: Literal['velocity', 'position']
+    velocity_goal: Optional[VelocityGoal] = None
+    position_goal: Optional[PositionGoal] = None

robogauge/tasks/pipeline/base_pipeline.py

@@ -10,15 +10,17 @@
 import traceback
 from robogauge.utils.logger import logger
 from robogauge.tasks.simulator import MujocoSimulator, MujocoConfig
-from robogauge.tasks.robots import BaseRobot, RobotConfig
+from robogauge.tasks.robots import BaseRobot, RobotConfig, Go2Config, Go2
 from robogauge.tasks.gauge import BaseGauge, BaseGaugeConfig
 
 class BasePipeline:
     def __init__(self, 
+        run_name: str,
         simulator_cfg: MujocoConfig,
         robot_cfg: RobotConfig,
         gauge_cfg: BaseGaugeConfig
     ):
+        self.run_name = run_name
         self.simulator_cfg = simulator_cfg
         self.robot_cfg = robot_cfg
         self.gauge_cfg = gauge_cfg
@@ -28,31 +30,34 @@ class BasePipeline:
         self.gauge: BaseGauge = eval(gauge_cfg.gauge_class)(gauge_cfg)
     
     def load(self):
+        logger.create_tensorboard(self.run_name)
         self.sim.load(
             self.gauge_cfg.assets.terrain_xml,
             self.robot_cfg.assets.robot_xml,
             self.gauge_cfg.assets.terrain_spawn_xy,
-            self.robot_cfg.assets.robot_spawn_height
+            self.robot_cfg.assets.robot_spawn_height,
+            self.robot_cfg.control.default_dof_pos
         )
     
     def run(self):
         try:
             self.load()
-            info = self.sim.step()
+            sim_data = self.sim.step()
             frame_skip = int(self.robot_cfg.control.control_dt / self.simulator_cfg.physics.simulation_dt)
             logger.info(f"Sim FPS: {1.0 / self.simulator_cfg.physics.simulation_dt:.2f}, Control FPS: {1.0 / self.robot_cfg.control.control_dt:.2f}, Frame Skip: {frame_skip:d}")
-            logger.info("Starting pipeline...")
+            logger.info("Running pipeline...")
             while not self.gauge.is_done():
                 goal = self.gauge.get_goal()
-                obs = self.robot.build_observation(info, goal)
+                obs = self.robot.build_observation(sim_data, goal)
-                action = self.robot.get_action(obs)
+                action, p_gains, d_gains, control_type = self.robot.get_action(obs)
+                self.sim.setup_action(action, p_gains, d_gains, control_type)
                 for _ in range(frame_skip):
-                    self.sim.apply_action(action)
+                    sim_data = self.sim.step()
-                    info = self.sim.step()
+                    self.gauge.update_metrics(sim_data)
-                    self.gauge.update_metrics(info)
                 if self.gauge.is_reset():
                     self.sim.reset()
-                    info = self.sim.step()
+                    sim_data = self.sim.step()
         finally:
             self.sim.close_viewer()
             logger.info("Pipeline execution finished.")
+            logger.info(f"Logging saved at: {logger.log_dir}")

robogauge/tasks/robots/__init__.py

@@ -1,4 +1,4 @@
 from .base_robot_config import RobotConfig
 from .base_robot import BaseRobot
-# from .go2.go2_config import Go2Config
+from .go2.go2_config import Go2Config
-# from .go2.go2_controller import Go2Controller
+from .go2.go2 import Go2

robogauge/tasks/robots/base_robot.py

@@ -9,21 +9,55 @@
 '''
 import torch
 import numpy as np
+
+from robogauge.utils.helpers import parse_path
+from robogauge.utils.logger import logger
 from robogauge.tasks.robots.base_robot_config import RobotConfig
+from robogauge.tasks.simulator.sim_data import SimData
+from robogauge.tasks.gauge.goal_data import GoalData
 
 class BaseRobot:
     def __init__(self, cfg: RobotConfig):
-        self.num_act = cfg.mdp.num_actions
+        self.cfg = cfg
-        self.num_obs = cfg.mdp.num_observations
+        self.device = self.cfg.control.device
-        self.model = None
+        self.num_obs = cfg.control.num_observations
+        self.num_action = cfg.control.num_actions
+        self.control_type = cfg.control.control_type
+        self.p_gains = np.array(cfg.control.p_gains)
+        self.d_gains = np.array(cfg.control.d_gains)
+        script_model_path = parse_path(cfg.control.torch_script_model_path)
+        logger.info(f"Loading robot model from '{script_model_path}'")
+        self.model = torch.jit.load(script_model_path).to(self.device)
+        self.model.eval()
     
-    def load_model(self):
+    def build_observation(self, sim_data: SimData, goal_data: GoalData) -> np.ndarray:
-        ...
+        obs = np.zeros(self.num_obs, dtype=np.float32)
-    
-    def build_observation(self, sim_info: dict, goal_info: dict) -> np.ndarray:
-        obs = np.zeros(self.num_obs)
         return obs
     
-    def get_action(self, obs) -> np.ndarray:
+    def get_action(self, obs: np.ndarray):
-        action = np.zeros_like(self.num_act)
+        """
-        return action
+        Returns:
+            action: (num_action,) target joint positions/velocities/torques
+            p_gains: (num_action,) proportional gains for Mujoco PD controller
+            d_gains: (num_action,) derivative gains for Mujoco PD controller
+            control_type: 'P', 'V', or 'T' for position/velocity/torque control
+        """
+        action = np.zeros(self.num_action, dtype=np.float32)
+        return action, self.p_gains, self.d_gains, self.control_type
+
+def get_projected_gravity(quat):
+    """ Compute world frame gravity (0, 0, -1) projected into robot base frame.
+    Args:
+        quat: (4,) quaternion (w, x, y, z) from robot base to world frame
+    Returns:
+        projected_gravity: (3,) projected gravity vector in robot base frame
+    """
+    qw, qx, qy, qz = quat
+
+    gravity_orientation = np.zeros(3)
+
+    gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
+    gravity_orientation[1] = -2 * (qz * qy + qw * qx)
+    gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
+
+    return gravity_orientation

robogauge/tasks/robots/base_robot_config.py

@@ -7,6 +7,7 @@
 @Blog    : https://wty-yy.github.io/
 @Desc    : Base Robot Configuration
 '''
+from typing_extensions import Literal
 from robogauge.utils.config import Config
 
 class RobotConfig(Config):
@@ -17,16 +18,32 @@ class RobotConfig(Config):
         robot_spawn_height = 0.1  # z [m]
 
     class control:
-        torch_script_model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_cts_61500.pt"
+        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  # target pos = action_scale * action * default_pos
+        control_type = 'P'  # Position control
-        stiffness = 20.0  # [N*m/rad]
-        damping = 0.5  # [N*m*s/rad]
 
-    class mdp:
+        # Mujoco joint PD gains
-        num_observations = 46
+        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:
+            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:
         lin_vel_x = [-1, 1]  # min max [m/s]
         lin_vel_y = [-1, 1]  # min max [m/s]

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]

robogauge/tasks/simulator/mujoco_config.py

@@ -13,7 +13,7 @@ class MujocoConfig(Config):
     simulator_class = 'MujocoSimulator'
 
     class physics:
-        simulation_dt = 0.005  # 200 Hz
+        simulation_dt = 0.002  # 500 Hz
     
     class viewer:
         headless = False
@@ -21,5 +21,6 @@ class MujocoConfig(Config):
     
     class render:
         save_video = False
+        video_fps = 30
         height = 480
         width = 640

robogauge/tasks/simulator/mujoco_simulator.py

@@ -5,7 +5,7 @@
 @Author  : wty-yy
 @Version : 1.0
 @Blog    : https://wty-yy.github.io/
-@Desc    : None
+@Desc    : Mujoco Simulator for Robogauge
 '''
 import mujoco
 import mujoco.viewer
@@ -15,11 +15,15 @@ import re
 import time
 import imageio
 import numpy as np
+from typing import Literal
 
 from robogauge.utils.logger import logger
-from robogauge.utils.helpers import pares_path
+from robogauge.utils.helpers import parse_path
 from robogauge.tasks.simulator.mujoco_config import MujocoConfig
-from robogauge.tasks.simulator.sim_data import RobotProprioception, JointState, BaseState, IMUState
+from robogauge.tasks.simulator.sim_data import (
+    SimData,
+    RobotProprioception, JointState, BaseState, IMUState
+)
 
 class MujocoSimulator:
     def __init__(self, sim_cfg: MujocoConfig):
@@ -28,12 +32,14 @@ class MujocoSimulator:
         self.robot_xml = None
         self.terrain_spawn_xy = None
         self.robot_spawn_height = None
+        self.default_dof_pos = None
         self.viewer = None
         self.renderer = None
         self.vid_writer = None
         self.vid_count = 0
         self._pause = True
         self.n_step = 0
+        self.sim_time = 0.0
 
     def load(
         self,
@@ -41,16 +47,19 @@ class MujocoSimulator:
         robot_xml: str = None,
         terrain_spawn_xy: list = None,
         robot_spawn_height: float = None,
+        default_dof_pos: list = None,
     ):
         """ Load terrain and robot into the simulator, support re-loading. """
         if terrain_xml is not None:
-            self.terrain_xml = pares_path(terrain_xml)
+            self.terrain_xml = parse_path(terrain_xml)
         if robot_xml is not None:
-            self.robot_xml = pares_path(robot_xml)
+            self.robot_xml = parse_path(robot_xml)
         if terrain_spawn_xy is not None:
             self.terrain_spawn_xy = terrain_spawn_xy
         if robot_spawn_height is not None:
             self.robot_spawn_height = robot_spawn_height
+        if default_dof_pos is not None:
+            self.default_dof_pos = default_dof_pos
 
         terrain_xml = self.terrain_xml
         robot_xml = self.robot_xml
@@ -58,6 +67,8 @@ class MujocoSimulator:
         robot_spawn_height = self.robot_spawn_height
         if terrain_xml is None or robot_xml is None:
             raise ValueError("Terrain and robot XML paths must be provided.")
+        if default_dof_pos is None:
+            raise ValueError("Default DOF positions must be provided.")
         
         robot_mjcf = mjcf.from_path(robot_xml)
         terrain_mjcf = mjcf.from_path(terrain_xml)
@@ -74,6 +85,10 @@ class MujocoSimulator:
         self.mj_model = self.mj_physics.model.ptr
         self.mj_data = self.mj_physics.data.ptr
         self.mj_model.opt.timestep = self.cfg.physics.simulation_dt
+        self.sim_dt = self.cfg.physics.simulation_dt
+        self.mj_data.qpos[7:] = default_dof_pos
+        mujoco.mj_forward(self.mj_model, self.mj_data)
+
         self.headless = self.cfg.viewer.headless
         if self.cfg.render.save_video and self.headless:
             logger.warning("Cannot save video in headless mode, disabling video saving.")
@@ -94,23 +109,33 @@ class MujocoSimulator:
                 vid_path = str(vid_dir / f"sim_video_{self.vid_count:03d}.mp4")
                 self.vid_writer = imageio.get_writer(
                     vid_path,
-                    fps=int(1 / self.cfg.physics.simulation_dt),
+                    fps=self.cfg.render.video_fps,
                 )
+                self.vid_frame_skip = int(1 / (self.cfg.render.video_fps * self.sim_dt * 2))
                 logger.info(f"Saving simulation video to: {vid_path}")
                 self.vid_count += 1
         self._pause = False
         self.n_step = 0
+        self.sim_time = 0.0
         self.preload_sensors()
 
+
+        # Robot controller placeholders
+        self.action = None
+        self.p_gains = None
+        self.d_gains = None
+        self.control_type = None
+
     def key_callback(self, keycode):
         if keycode == 32:
             self._pause = not self._pause
             logger.info(f"Pause toggled: {self._pause}")
     
-    def step(self) -> dict:
+    def step(self) -> SimData:
         """ Simulation step, pause will block thread. """
         while self._pause:
             time.sleep(0.1)
+        self.update_torque()
         self.mj_physics.step()
         if self.viewer is not None:
             if self.viewer.is_running():
@@ -120,7 +145,7 @@ class MujocoSimulator:
                 if time_untile_next_render > 0:
                     time.sleep(time_untile_next_render)
                 self.viewer.sync()
-                if self.vid_writer is not None:
+                if self.vid_writer is not None and self.n_step % self.vid_frame_skip == 0:
                     self.renderer.update_scene(self.mj_data, camera=self.viewer.cam)
                     frame = self.renderer.render()
                     self.vid_writer.append_data(frame)
@@ -129,43 +154,75 @@ class MujocoSimulator:
                 logger.warning("Viewer closed by user, stop video recording.")
                 self.close_viewer()
 
-        n_sensor = self.mj_model.nsensor
+        self.proprio = proprio = RobotProprioception(
-        
-        proprio = RobotProprioception(
             joint=JointState(
                 pos=self.get_sensor_data('joint_pos'),
                 vel=self.get_sensor_data('joint_vel'),
                 force=self.get_sensor_data('joint_eff'),
             ),
             imu=IMUState(
-                quat=self.get_sensor_data('imu_quat'),
-                ang_vel=self.get_sensor_data('imu_ang_vel'),
-                acc=self.get_sensor_data('imu_acc'),
                 pos=self.get_sensor_data('imu_pos'),
+                quat=self.get_sensor_data('imu_quat'),
+                acc=self.get_sensor_data('imu_acc'),
                 lin_vel=self.get_sensor_data('imu_lin_vel'),
+                ang_vel=self.get_sensor_data('imu_ang_vel'),
             ),
             base=BaseState(
                 pos=self.mj_data.qpos[:3],      # world frame
                 quat=self.mj_data.qpos[3:7],    # world frame
-                vel=self.mj_data.qvel[:3],      # body frame
+                lin_vel=self.mj_data.qvel[:3],  # body frame
                 ang_vel=self.mj_data.qvel[3:6], # body frame
             )
         )
+        if self.n_step % int(0.1 / self.sim_dt) == 0:
+            logger.log(value=np.mean(proprio.imu.quat - proprio.base.quat), tag="sim/delta_quat", step=self.n_step)
+            logger.log(value=np.mean(proprio.imu.ang_vel - proprio.base.ang_vel), tag="sim/delta_ang_vel", step=self.n_step)
+            logger.log(value=np.mean(proprio.imu.lin_vel - proprio.base.lin_vel), tag="sim/delta_lin_vel", step=self.n_step)
         if self.n_step == 0:
-            self.debug_print_proprio_shapes(proprio)
+            self.debug_print_proprio_shapes()
-        self.n_step += 1
 
-        return proprio
+        sim_data = SimData(
+            n_step=self.n_step,
+            sim_dt=self.sim_dt,
+            proprio=proprio
+        )
+
+        self.n_step += 1
+        self.sim_time = self.n_step * self.sim_dt
+        return sim_data
     
     def reset(self):
         """ Reset the simulator to initial state. """
         self.mj_physics.reset()
+        self.mj_data.qpos[7:] = self.default_dof_pos
+        mujoco.mj_forward(self.mj_model, self.mj_data)
+
+        self.action = None
         if self.viewer is not None:
             self.viewer.sync()
 
-    def apply_action(self, action: np.ndarray):
+    def setup_action(self,
-        """ Apply action to the simulator. """
+            action: np.ndarray,
-        self.mj_data.ctrl[:] = action
+            p_gains: np.ndarray = None,
+            d_gains: np.ndarray = None,
+            control_type: Literal['P'] = 'P'
+        ):
+        """ Setup action to the simulator. """
+        self.action = action
+        self.p_gains = p_gains
+        self.d_gains = d_gains
+        self.control_type = control_type
+    
+    def update_torque(self):
+        if self.action is None:
+            return
+        dof_pos = self.proprio.joint.pos
+        dof_vel = self.proprio.joint.vel
+        if self.control_type == 'P':
+            torques = self.p_gains * (self.action - dof_pos) - self.d_gains * dof_vel
+        else:
+            raise NotImplementedError(f"Control type '{self.control_type}' not implemented.")
+        self.mj_data.ctrl[:] = torques
     
     def close_viewer(self):
         """ Close the viewer and video writer. """
@@ -189,12 +246,16 @@ class MujocoSimulator:
         self.imu_acc = self.find_sensors(tag_name="accelerometer")
         self.imu_pos = self.find_sensors(tag_name="framepos")
         self.imu_lin_vel = self.find_sensors(tag_name="framelinvel")
+        actuator_names = [mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) for i in range(self.mj_model.nu)]
         logger.info(
             f"\n{'='*20} XML SENSOR NAMES {'='*20}\n"
-            f"""Joint Position Sensors [{len(self.joint_pos_sensor_names)}]: {self.joint_pos_sensor_names}\n"""
+            f"""Joint Position Sensors [{len(self.joint_pos_sensor_names)}]: {[x.rsplit('/')[-1] for x in self.joint_pos_sensor_names]}\n"""
-            f"""Joint Velocity Sensors [{len(self.joint_vel_sensor_names)}]: {self.joint_vel_sensor_names}\n"""
+            f"""Joint Velocity Sensors [{len(self.joint_vel_sensor_names)}]: {[x.rsplit('/')[-1] for x in self.joint_vel_sensor_names]}\n"""
-            f"""Joint Effort Sensors [{len(self.joint_eff_sensor_names)}]: {self.joint_eff_sensor_names}\n"""
+            f"""Joint Effort Sensors [{len(self.joint_eff_sensor_names)}]: {[x.rsplit('/')[-1] for x in self.joint_eff_sensor_names]}\n"""
-            f"""IMU Sensors: Quat{self.imu_quat}, AngVel{self.imu_ang_vel}, Acc{self.imu_acc}, Pos{self.imu_pos}, LinVel{self.imu_lin_vel}"""
+            f"""Actuators [{len(actuator_names)}]: {[x.rsplit('/')[-1] for x in actuator_names]}\n"""
+            f"""IMU Sensors: Quat{self.imu_quat}, AngVel{self.imu_ang_vel}, Acc{self.imu_acc}, Pos{self.imu_pos}, LinVel{self.imu_lin_vel}\n"""
+            f"""!!!Checkout actuators order is consistent with joint sensors!!!\n"""
+            f"{'='*58}"
         )
 
         # Cache sensor indices
@@ -259,7 +320,7 @@ class MujocoSimulator:
             data_list.append(self.mj_data.sensordata[adr:adr+dim])
         return np.concatenate(data_list)
 
-    def debug_print_proprio_shapes(self, proprio: RobotProprioception):
+    def debug_print_proprio_shapes(self):
         """Log shapes (or lengths) of each numpy vector inside a RobotProprioception.
 
         This helps debug mismatched sensor sizes between robots.
@@ -271,9 +332,9 @@ class MujocoSimulator:
             except Exception:
                 return None
 
-        jp = proprio.joint
+        jp = self.proprio.joint
-        bs = proprio.base
+        bs = self.proprio.base
-        imu = proprio.imu
+        imu = self.proprio.imu
 
         logger.info("Proprioception shapes:")
         logger.info(f"  joint.pos: { _shape(jp.pos) }")
@@ -282,7 +343,7 @@ class MujocoSimulator:
 
         logger.info(f"  base.pos: { _shape(bs.pos) }")
         logger.info(f"  base.quat: { _shape(bs.quat) }")
-        logger.info(f"  base.vel: { _shape(bs.vel) }")
+        logger.info(f"  base.vel: { _shape(bs.lin_vel) }")
         logger.info(f"  base.ang_vel: { _shape(bs.ang_vel) }")
 
         logger.info(f"  imu.quat: { _shape(imu.quat) }")

robogauge/tasks/simulator/sim_data.py

@@ -11,19 +11,25 @@ class JointState:
 class BaseState:
     pos: np.ndarray
     quat: np.ndarray
-    vel: np.ndarray
+    lin_vel: np.ndarray
     ang_vel: np.ndarray
 
 @dataclass
 class IMUState:
-    quat: np.ndarray
-    ang_vel: np.ndarray
-    acc: np.ndarray
     pos: np.ndarray
+    quat: np.ndarray
+    acc: np.ndarray
     lin_vel: np.ndarray
+    ang_vel: np.ndarray
 
 @dataclass
 class RobotProprioception:
     joint: JointState
     base: BaseState
     imu: IMUState
+
+@dataclass
+class SimData:
+    n_step: int
+    sim_dt: float
+    proprio: RobotProprioception

robogauge/utils/helpers.py

@@ -14,7 +14,7 @@ from argparse import ArgumentParser
 from pathlib import Path
 from robogauge import ROBOGAUGE_ROOT_DIR
 
-def pares_path(path):
+def parse_path(path):
     if "{ROBOGAUGE_ROOT_DIR}" in str(path):
         path = str(path).replace("{ROBOGAUGE_ROOT_DIR}", ROBOGAUGE_ROOT_DIR)
     return path
@@ -47,10 +47,13 @@ def parse_args():
     parser = ArgumentParser()
     parameters = [
         {"name": "--task-name", "type": str, "default": "base", "help": "Name of the task to run."},
-        {"name": "--experiment-name", "type": str, "default": "base", "help": "Name of the experiment to run."},
+        {"name": "--experiment-name", "type": str, "help": "Name of the experiment to run."},
         {"name": "--headless", "action": "store_true", "default": False, "help": "Run in headless mode."},
         {"name": "--save-video", "action": "store_true", "default": False, "help": "Save video output."},
     ]
     for param in parameters:
         parser.add_argument(param['name'], **{k: v for k, v in param.items() if k != 'name'})
-    return parser.parse_args()
+    args = parser.parse_args()
+    if args.experiment_name is None:
+        args.experiment_name = f"{args.task_name}_exp"
+    return args

robogauge/utils/logger.py

@@ -14,6 +14,7 @@ import time
 import logging
 from pathlib import Path
 from robogauge import ROBOGAUGE_ROOT_DIR
+from torch.utils.tensorboard import SummaryWriter
 
 class LogColor:
     """ ANSI color codes """
@@ -49,6 +50,10 @@ class ColorFormatter(logging.Formatter):
 
 
 class Logger:
+    logger: logging.Logger = None
+    log_dir: Path = None
+    writer: SummaryWriter = None
+
     def create(self,
         experiment_name,
         console_output=True, color_output=True,
@@ -94,6 +99,12 @@ class Logger:
             fh.setFormatter(file_formatter)
             self.logger.addHandler(fh)
     
+    def create_tensorboard(self, run_name: str):
+        if self.writer is not None:
+            self.writer.close()
+        self.writer = SummaryWriter(str(self.log_dir / run_name))
+        self.info(f"Tensorboard writer created at: {self.log_dir / run_name}")
+
     def debug(self, msg, *args, **kwargs):
         self.logger.debug(msg, *args, **kwargs, stacklevel=2)
     
@@ -109,6 +120,19 @@ class Logger:
     def critical(self, msg, *args, **kwargs):
         self.logger.critical(msg, *args, **kwargs, stacklevel=2)
     
+    def log(self, value, tag, step):
+        """ Log scalar value to tensorboard
+
+        Args:
+            value (float): scalar value
+            tag (str): tag name
+            step (int): step number
+        """
+        if self.writer is not None:
+            self.writer.add_scalar(tag, value, step)
+        else:
+            self.warning("Tensorboard writer is not initialized, skipping log.")
+
 logger = Logger()
 
 if __name__ == '__main__':

robogauge/utils/task_register.py

@@ -33,7 +33,7 @@ class TaskRegister():
         robot_cfg = self.robot_cfgs[name]
         return sim_cfg, gauger_cfg, robot_cfg
     
-    def make_pipeline(self, name, args=None, sim_cfg=None, gauger_cfg=None, robot_cfg=None):
+    def make_pipeline(self, name, args=None, sim_cfg=None, gauger_cfg=None, robot_cfg=None, run_name='0'):
         default_cfgs = self.get_cfgs(name)
         if sim_cfg is None:
             sim_cfg = default_cfgs[0]
@@ -44,7 +44,7 @@ class TaskRegister():
         if args is not None:
             self.update_args_to_cfg(sim_cfg, gauger_cfg, robot_cfg, args)
         pipeline_class = self.get_pipeline_class(name)
-        return pipeline_class(sim_cfg, robot_cfg, gauger_cfg)
+        return pipeline_class(run_name, sim_cfg, robot_cfg, gauger_cfg)
 
     def update_args_to_cfg(self, sim_cfg, gauger_cfg, robot_cfg, args):
         if args.headless is not None: