v1.1.1-rc2; add velocity visualization

This commit is contained in:
wty-yy
2026-01-26 10:43:48 +08:00
parent 9b026c9e7a
commit 493fa04341
5 changed files with 119 additions and 32 deletions

View File

@@ -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. 加入可视化潜空间的绘图工具,使用方法:

Binary file not shown.

View File

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

View File

@@ -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):

View File

@@ -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.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.update_external_rendering(self.viewer, ctype='viewer')
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)
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]
)
self.update_external_rendering(self.renderer, ctype='renderer')
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