v1.1.1-rc2; add velocity visualization

UPDATE.md

@@ -1,4 +1,9 @@
 # UPDATE
+## 20260126
+### v1.1.1-rc2
+1. 加入可视化当前速度（蓝），目标速度（绿）箭头
+2. moe默认模型换为`go2_moe_cts_137k_0.6745.pt`
+3. 修复1080p在某些地形无法渲染的问题，将[mujoco_config.py](robogauge/tasks/simulator/mujoco_config.py)中的`render.width, render.height`改为1920x1080即可使用，默认640x480
 ## 20260122
 ### v1.1.1-rc1
 1. 加入可视化潜空间的绘图工具，使用方法：

robogauge/tasks/pipeline/base_pipeline.py

@@ -92,10 +92,9 @@ class BasePipeline:
                     sim_data = self.reset_sim_and_robot(sim_data)
                     continue
 
-                if goal_data.visualization_pos is not None:
+                # Setup visualization for simulation
-                    self.sim.set_target_pos(goal_data.visualization_pos)
+                self.sim.target_pos = goal_data.visualization_pos
-                else:
+                self.sim.target_velocity = goal_data.velocity_goal
-                    self.sim.set_target_pos(None)
 
                 obs = self.robot.build_observation(self.add_noise(sim_data), goal_data)
                 action, p_gains, d_gains, control_type = self.robot.get_action(obs)

robogauge/tasks/robots/go2/go2_moe_config.py

@@ -13,7 +13,7 @@ class Go2MoEConfig(Go2Config):
     robot_class = 'Go2MoE'
 
     class control(Go2Config.control):
-        model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_moe_cts_124k.pt"
+        model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go2/go2_moe_cts_137k_0.6745.pt"
         save_additional_output = False
 
 class Go2MoETerrainConfig(Go2MoEConfig):

robogauge/tasks/simulator/mujoco_simulator.py

@@ -16,7 +16,7 @@ import time
 import imageio
 import numpy as np
 from pathlib import Path
-from typing import Literal, List
+from typing import Literal, List, Optional, Union
 
 from robogauge.utils.logger import logger
 from robogauge.utils.helpers import parse_path
@@ -26,6 +26,7 @@ from robogauge.tasks.simulator.sim_data import (
     SimData,
     RobotProprioception, JointState, BaseState, IMUState
 )
+from robogauge.tasks.gauge.goal_data import VelocityGoal
 
 class MujocoSimulator:
     def __init__(self, sim_cfg: MujocoConfig):
@@ -45,6 +46,7 @@ class MujocoSimulator:
         self.n_step = 0
         self.sim_time = 0.0
         self.target_pos = None
+        self.target_velocity: Optional[VelocityGoal] = None
         self.penetration_reset_count = 0
 
     def load(
@@ -78,6 +80,11 @@ class MujocoSimulator:
         # Create MJCF models
         robot_mjcf = mjcf.from_path(robot_xml)
         terrain_mjcf = mjcf.from_path(terrain_xmls[0])
+        visual_elem = terrain_mjcf.visual
+        global_elem = visual_elem.get_children('global')
+        global_elem.offwidth = 1920
+        global_elem.offheight = 1080
+
         for path in terrain_xmls[1:]:
             next_terrain = mjcf.from_path(path)
             terrain_mjcf.attach(next_terrain)
@@ -180,9 +187,6 @@ class MujocoSimulator:
         if keycode == 32:
             self._pause = not self._pause
             logger.info(f"Pause toggled: {self._pause}")
-    
-    def set_target_pos(self, pos):
-        self.target_pos = pos
 
     def step(self) -> SimData:
         """ Simulation step, pause will block thread. """
@@ -194,17 +198,7 @@ class MujocoSimulator:
         # Viewer sync
         if self.viewer is not None:
             if self.viewer.is_running():
-                if self.target_pos is not None:
+                self.update_external_rendering(self.viewer, ctype='viewer')
-                    self.viewer.user_scn.ngeom = 0
-                    mujoco.mjv_initGeom(
-                        self.viewer.user_scn.geoms[0],
-                        type=mujoco.mjtGeom.mjGEOM_SPHERE,
-                        size=[0.1, 0, 0],
-                        pos=self.target_pos,
-                        mat=np.eye(3).flatten(),
-                        rgba=[1, 0, 0, 1]
-                    )
-                    self.viewer.user_scn.ngeom = 1
                 self.viewer.sync()
                 time_untile_next_render = self.cfg.physics.simulation_dt - (
                     time.time() - self.last_render_time
@@ -221,18 +215,7 @@ class MujocoSimulator:
             render_cam = self.viewer.cam if self.viewer is not None else self.offscreen_cam
             # mujoco.mjv_updateCamera(render_cam)
             self.renderer.update_scene(self.mj_data, camera=render_cam)
-            
+            self.update_external_rendering(self.renderer, ctype='renderer')
-            if self.target_pos is not None:
-                self.renderer.scene.ngeom += 1
-                mujoco.mjv_initGeom(
-                    self.renderer.scene.geoms[self.renderer.scene.ngeom - 1],
-                    type=mujoco.mjtGeom.mjGEOM_SPHERE,
-                    size=[0.1, 0, 0],
-                    pos=self.target_pos,
-                    mat=np.eye(3).flatten(),
-                    rgba=[1, 0, 0, 1]
-                )
-
             frame = self.renderer.render()
             self.vid_writer.append_data(frame)
 
@@ -282,6 +265,106 @@ class MujocoSimulator:
         self.check_truncation(sim_data)
         return sim_data
     
+    def update_external_rendering(self,
+            handle: Union[mujoco.viewer.Handle, mujoco.Renderer],
+            ctype: Literal['viewer', 'renderer'],
+        ):
+        """ Update external rendering handle (viewer or renderer). """
+        def add_target_sphere(geom_elem):
+            mujoco.mjv_initGeom(
+                geom_elem,
+                type=mujoco.mjtGeom.mjGEOM_SPHERE,
+                size=[0.1, 0, 0],
+                pos=self.target_pos,
+                mat=np.eye(3).flatten(),
+                rgba=[1, 0, 0, 1]
+            )
+
+        def add_arrow(geom_elem, pos, vec, rgba, scale=1.0):
+            vel_norm = np.linalg.norm(vec)
+            if vel_norm < 1e-3: 
+                mujoco.mjv_initGeom(
+                    geom_elem,
+                    type=mujoco.mjtGeom.mjGEOM_NONE,
+                    size=[0, 0, 0], pos=pos, mat=np.eye(3).flatten(), rgba=[0, 0, 0, 0]
+                )
+                return
+
+            mat = np.zeros(9)
+            target_quat = np.zeros(4)
+            vec_normalized = vec / vel_norm
+            mujoco.mju_quatZ2Vec(target_quat, vec_normalized)
+            mujoco.mju_quat2Mat(mat, target_quat)
+            total_length = max(vel_norm * scale * 0.5, 0.05)
+            shaft_radius = 0.015
+            head_radius = 0.035
+            mujoco.mjv_initGeom(
+                geom_elem,
+                type=mujoco.mjtGeom.mjGEOM_ARROW,
+                size=[shaft_radius, head_radius, total_length], 
+                pos=pos,
+                mat=mat,
+                rgba=rgba
+            )
+
+        viewer_geom_idx = 0
+        if ctype == 'viewer':
+            handle.user_scn.ngeom = 0  # reset user scene geometry
+
+        if self.target_pos is not None:
+            if ctype == 'viewer':
+                add_target_sphere(handle.user_scn.geoms[viewer_geom_idx])
+                viewer_geom_idx += 1
+            else:
+                handle.scene.ngeom += 1
+                add_target_sphere(handle.scene.geoms[self.renderer.scene.ngeom - 1])
+        
+        if self.target_velocity is not None:
+            base_pos_world = self.mj_data.qpos[:3]
+            base_quat = self.mj_data.qpos[3:7]
+            
+            # rendering arrows start position
+            offset_body = np.array([0.0, 0.0, 0.6])
+            offset_world = np.zeros(3)
+            mujoco.mju_rotVecQuat(offset_world, offset_body, base_quat)
+            start_pos = base_pos_world + offset_world
+
+            # Body Frame -> World Frame
+            target_lin_vel_body = np.array([
+                self.target_velocity.lin_vel_x, 
+                self.target_velocity.lin_vel_y, 
+                0.0
+            ])
+            target_lin_vel_world = np.zeros(3)
+            mujoco.mju_rotVecQuat(target_lin_vel_world, target_lin_vel_body, base_quat)
+
+            # Body Frame -> World Frame
+            raw_current_vel = self.proprio.base.lin_vel
+            current_lin_vel_body = np.array([
+                raw_current_vel[0],
+                raw_current_vel[1],
+                0.0
+            ])
+            current_lin_vel_world = np.zeros(3)
+            mujoco.mju_rotVecQuat(current_lin_vel_world, current_lin_vel_body, base_quat)
+
+            COLOR_GREEN = [0, 1, 0, 1] # target
+            COLOR_BLUE = [0, 0, 1, 1]  # current
+
+            if ctype == 'viewer':
+                add_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, target_lin_vel_world, COLOR_GREEN)
+                viewer_geom_idx += 1
+                add_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, current_lin_vel_world, COLOR_BLUE)
+                viewer_geom_idx += 1
+            else:
+                handle.scene.ngeom += 1
+                add_arrow(handle.scene.geoms[handle.scene.ngeom - 1], start_pos, target_lin_vel_world, COLOR_GREEN)
+                handle.scene.ngeom += 1
+                add_arrow(handle.scene.geoms[handle.scene.ngeom - 1], start_pos, current_lin_vel_world, COLOR_BLUE)
+
+        if ctype == 'viewer':
+            handle.user_scn.ngeom = viewer_geom_idx
+    
     def check_penetration(self, threshold: float = -0.02):
         if self.penetration_reset_count >= self.cfg.truncation.penetration_max_reset_num:
             return False, None, None, None