cyy/RoboGauge
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

v1.1.4; Add ZMP margin metric

Browse Source
This commit is contained in:
wty-yy
2026-03-23 11:24:00 +08:00
parent be054d02b0
commit fe35afea22
17 changed files with 782 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@@ -11,3 +11,4 @@ logs/
logs*/
.aim/
*.bkp
.vscode/

10
CMD.md
View File

@@ -18,7 +18,7 @@ python robogauge/scripts/run.py \
--headless
# To evaluate metrics for level terrains, add `--goals` and `--spawn-type`
# level terrains: wave, slope_fb, slope_bd, stairs_fb, stairs_bd, obstacle (default is target_pos)
# level terrains: wave, slope_fd, slope_bd, stairs_fd, stairs_bd, obstacle (default is target_pos)
python robogauge/scripts/run.py \
--task go2_moe.obstacle \
--experiment-name debug \
@@ -26,6 +26,14 @@ python robogauge/scripts/run.py \
--friction 2 \
--spawn-type level_eval \
--goals max_velocity diagonal_velocity
# For joystick control and tensorboard DEBUG
python robogauge/scripts/run.py \
--task go2_moe.flat \
--experiment-name debug \
--goals joystick \
--write-tensorboard \
--save-video
```
# Multi Pipeline

1
README.md
View File

@@ -133,6 +133,7 @@ Supported per-step metrics (measured after each `env.step`). **All metrics are n
| 4 | `dof_power` | Motor energy consumption | Scaling factor | 100 | `1-x` |
| 5 | `orientation_stability` | Body orientation stability (Roll) | NA | NA | `1-x` |
| 6 | `torque_smoothness` | Torque smoothness | Scaling factor | 30 | `1-x` |
| 7 | `zmp_margin` | ZMP margin | NA | Default diagonal foot distance | `1-x` |
### Velocity-Tracking Targets

1
README_zh.md
View File

@@ -119,6 +119,7 @@ while True:
| 4 | `dof_power` | 电机耗能 | 缩放系数 | 100 | `1-x` |
| 5 | `orientation_stability` | 机身姿态稳定性 (Roll) | NA | NA | `1-x` |
| 6 | `torque_smoothness` | 力矩平滑度 | 缩放系数 | 30 | `1-x` |
| 7 | `zmp_margin` | ZMP裕度 | NA | 默认状态下对角足端间距 | `1-x` |
### 速度追踪目标
针对在虚实迁移中发现的问题, 整理指标 (metrics) 内容如下:

4
UPDATE.md
View File

@@ -1,4 +1,8 @@
# UPDATE
## 20260322
### v1.1.4
1. 添加ZMP裕度指标
2. 将默认评估速度统一为1m/s
## 20260210
### v1.1.3
1. 修复摩擦计算问题，之前的所有stress的摩擦*0.4后才是真实摩擦系数，所以之前用的摩擦范围是0.2~1.0，步长0.1

79
assets/docs/zmp_margin.md Normal file
View File

@@ -0,0 +1,79 @@
# zmp_margin
参考论文：[ZERO-MOMENT POINT -- THIRTY FIVE YEARS OF ITS LIFE](https://www.worldscientific.com/doi/abs/10.1142/S0219843604000083)
## 基本定义
- Zero Moment Point（ZMP，零矩点）：机械系统的惯性力和重力的净力矩在该点沿水平轴分量为零。
- Support Polygon（支撑多边形）：机械系统所有与地面接触点构成的凸包。
- FZMP（fictitious ZMP，虚拟零矩点）：当 ZMP 落在支撑多边形外时，称 ZMP 为 FZMP，此时系统可能处于非平衡状态。
- Virtual Horizontal Plane（虚拟水平投影面）：机械系统所有与地面接触点的几何中心记为 $O'$，将世界坐标系平移到 $O'$，得到的 $xy$ 平面。
## MuJoCo 中的符号约定
设机器人总共包含 $N$ 个刚体。对于第 $i$ 个刚体，在当前时刻具有如下量：
- 质量：$m_i$
- 质心位置：$\boldsymbol{p}_i$
这里相对于当前 `base_link` 在世界坐标系 $z$ 轴方向上平移到虚拟水平投影面的点。
- 质心线加速度：$\boldsymbol{\ddot{p}}_i$
- 角速度：$\boldsymbol{\omega}_i$
- 角加速度：$\boldsymbol{\dot{\omega}}_i$
- 惯性张量：$\boldsymbol{I}_i$
上述量均相对于世界坐标系，保持相对于惯性参考系表达；本体坐标系不一定是参考系。计算 $\boldsymbol{I}_i$ 时需要确保其已经转换到世界坐标系下。
## 总力与总力矩
整个系统的总力与总力矩定义为：
$$
\boldsymbol{F}_{\text{total}} = \sum_{i=1}^{N} m_i \left(\boldsymbol{g} - \boldsymbol{\ddot{p}}_i\right)
$$
$$
\boldsymbol{M}_{\text{total}} =
\sum_{i=1}^{N}
\left[
\left(\boldsymbol{p}_i \times m_i \left(\boldsymbol{g} - \boldsymbol{\ddot{p}}_i\right)\right)
- \left(\boldsymbol{I}_i \boldsymbol{\dot{\omega}}_i + \boldsymbol{\omega}_i \times \left(\boldsymbol{I}_i \boldsymbol{\omega}_i\right)\right)
\right]
$$
## ZMP 坐标
由
$$
\boldsymbol{M} = \boldsymbol{r} \times \boldsymbol{F}
$$
可得
$$
\begin{cases}
\boldsymbol{M}_y = -x_{\text{zmp}} \boldsymbol{F}_z + z_{\text{zmp}} \boldsymbol{F}_x \\
\boldsymbol{M}_x = \phantom{-}y_{\text{zmp}} \boldsymbol{F}_z + z_{\text{zmp}} \boldsymbol{F}_y
\end{cases}
$$
当 $z_{\text{zmp}} = 0$ 时，有
$$
\begin{cases}
x_{\text{zmp}} = -\dfrac{\boldsymbol{M}_y}{\boldsymbol{F}_z} \\
y_{\text{zmp}} = \phantom{-}\dfrac{\boldsymbol{M}_x}{\boldsymbol{F}_z}
\end{cases}
$$
## 最终指标
令默认状态下的对角足端距离为 $D_{\text{norm}}$，则总的 `zmp_margin` 定义为：
$$
m_{\text{zmp_margin}} =
\max \left(
0,
1 - \dfrac{\left\| (x_{\text{zmp}}, y_{\text{zmp}}) \right\|_2}{D_{\text{norm}}}
\right)
$$

3
robogauge/tasks/gauge/base_gauge.py
View File

@@ -138,13 +138,14 @@ class BaseGauge:
def update_metrics(self, sim_data: SimData, goal_data: GoalData):
if sim_data.n_step % int(self.cfg.metrics.metric_dt / sim_data.sim_dt + 1e-9) != 0:
return
return False
metrics_results = {}
for metric_name, metric_obj in zip(self.info['metric'], self.metrics):
val = metric_obj(sim_data, goal_data)
if metric_name not in ['visualization']:
metrics_results[metric_name] = val
self.goals[self.goal_idx].update_metrics(metrics_results)
return True
def reset_metrics(self):
for metric in self.metrics:

10
robogauge/tasks/gauge/base_gauge_config.py
View File

@@ -16,6 +16,7 @@ QUALITY_WEIGHTS = { # Weights for geometric average, to calculate quality scor
'dof_power': 1,
'orientation_stability': 1,
'torque_smoothness': 1,
'zmp_margin': 1,
}
class BaseGaugeConfig(Config):
@@ -81,3 +82,12 @@ class BaseGaugeConfig(Config):
class torque_smoothness:
enabled = True
scaling_factor = 30.0 # [Nm] scaling factor for torque smoothness metric
class zmp_margin:
enabled = True
contact_threshold = 1e-3 # [m] contact.dist <= threshold is treated as support contact
force_threshold = 1e-6 # [N] abs(Fz) below threshold skips ZMP and returns 1.0
draw_point = True # Whether to draw the ZMP point in both viewer and offscreen render
draw_point_size = 0.03 # [m] ZMP point sphere radius in Mujoco visualization
draw_height_offset = 0.00 # [m] Lift the point slightly above the support plane for visibility
draw_point_rgba = [1.0, 0.85, 0.1, 1.0] # RGBA color of the ZMP point

7
robogauge/tasks/gauge/goals/base_goal.py
View File

@@ -49,10 +49,15 @@ class BaseGoal:
def update_metrics(self, metrics: dict):
""" Update step metrics for the current goal."""
quality_score = 1.0
quality_weight_sum = 0.0
for metric_name, value in metrics.items():
self.goal_metrics[metric_name].append(value)
if metric_name not in QUALITY_WEIGHTS:
continue
quality_score *= min(max(1e-9, value), 1.0) ** QUALITY_WEIGHTS[metric_name]
quality_score = quality_score ** (1.0 / sum(QUALITY_WEIGHTS.values()))
quality_weight_sum += QUALITY_WEIGHTS[metric_name]
if quality_weight_sum > 0.0:
quality_score = quality_score ** (1.0 / quality_weight_sum)
self.goal_quality_scores.append(quality_score)
@property

2
robogauge/tasks/gauge/metrics/__init__.py
View File

@@ -2,4 +2,4 @@ from .base_metric import BaseMetric
from .dof_metrics import DofLimitsMetric, DofPowerMetric
from .visualization import VisualizationMetric
from .vel_metrics import LinVelErrMetric, AngVelErrMetric
from .stable_metric import OrientationStabilityMetric, TorqueSmoothnessMetric
from .stable_metric import OrientationStabilityMetric, TorqueSmoothnessMetric, ZmpMarginMetric

213
robogauge/tasks/gauge/metrics/stable_metric.py
View File

@@ -1,11 +1,11 @@
# -*- coding: utf-8 -*-
'''
@File : height_metric.py
@File : stable_metric.py
@Time : 2025/12/18 20:18:33
@Author : wty-yy
@Version : 1.0
@Blog : https://wty-yy.github.io/
@Desc : Orientation Stability, Torque Smoothness Metric Implementation
@Desc : Stability-related metrics implementation
'''
import numpy as np
@@ -17,8 +17,8 @@ from robogauge.utils.logger import logger
class OrientationStabilityMetric(BaseMetric):
""" Metric to log height stability. """
name = 'height_std_metric'
""" Metric to log body orientation stability. """
name = 'orientation_stability_metric'
def __init__(self,
robot_cfg: RobotConfig,
@@ -60,3 +60,208 @@ class TorqueSmoothnessMetric(BaseMetric):
metric_value = 1.0 - rms_value / self.scaling_factor
logger.log(rms_value, f'stable_metric/torque_rms_diff', step=sim_data.n_step)
return metric_value
def _cross_2d(o: np.ndarray, a: np.ndarray, b: np.ndarray) -> float:
oa = a - o
ob = b - o
return float(oa[0] * ob[1] - oa[1] * ob[0])
def _convex_hull_2d(points: np.ndarray) -> np.ndarray:
if points.shape[0] <= 1:
return points.copy()
pts = np.unique(points, axis=0)
if pts.shape[0] <= 1:
return pts
pts = pts[np.lexsort((pts[:, 1], pts[:, 0]))]
lower = []
for point in pts:
while len(lower) >= 2 and _cross_2d(lower[-2], lower[-1], point) <= 0.0:
lower.pop()
lower.append(point)
upper = []
for point in pts[::-1]:
while len(upper) >= 2 and _cross_2d(upper[-2], upper[-1], point) <= 0.0:
upper.pop()
upper.append(point)
hull = np.array(lower[:-1] + upper[:-1], dtype=np.float32)
return hull if hull.size > 0 else pts[:1]
def _polygon_area_2d(points: np.ndarray) -> float:
if points.shape[0] < 3:
return 0.0
x = points[:, 0]
y = points[:, 1]
return 0.5 * abs(float(np.dot(x, np.roll(y, -1)) - np.dot(y, np.roll(x, -1))))
def _point_on_segment_2d(point: np.ndarray, start: np.ndarray, end: np.ndarray) -> bool:
seg = end - start
rel = point - start
seg_norm_sq = float(np.dot(seg, seg))
if seg_norm_sq == 0.0:
return float(np.linalg.norm(rel)) == 0.0
cross = abs(seg[0] * rel[1] - seg[1] * rel[0])
if cross > 0.0:
return False
dot = float(np.dot(rel, seg))
return 0.0 <= dot <= seg_norm_sq
def _point_in_support_region(point: np.ndarray, support_polygon: np.ndarray) -> bool:
if support_polygon.shape[0] == 0:
return False
if support_polygon.shape[0] == 1:
return float(np.linalg.norm(point - support_polygon[0])) == 0.0
if support_polygon.shape[0] == 2:
return _point_on_segment_2d(point, support_polygon[0], support_polygon[1])
prev_sign = None
for idx in range(support_polygon.shape[0]):
start = support_polygon[idx]
end = support_polygon[(idx + 1) % support_polygon.shape[0]]
cross = _cross_2d(start, end, point)
if cross == 0.0:
continue
cur_sign = cross > 0.0
if prev_sign is None:
prev_sign = cur_sign
elif prev_sign != cur_sign:
return False
return True
class ZmpMarginMetric(BaseMetric):
""" Zero Moment Point metric. """
name = 'zmp_margin_metric'
D_NORM_MIN = 1e-5
def __init__(self,
robot_cfg: RobotConfig,
contact_threshold: float = 1e-3,
force_threshold: float = 1e-6,
draw_point: bool = False,
draw_point_size: float = 0.03,
draw_height_offset: float = 0.02,
draw_point_rgba: list = None,
**kwargs
):
super().__init__(robot_cfg)
self.contact_threshold = contact_threshold
self.force_threshold = force_threshold
self.draw_point = draw_point
self.draw_point_size = draw_point_size
self.draw_height_offset = draw_height_offset
if draw_point_rgba is None:
draw_point_rgba = [1.0, 0.85, 0.1, 1.0]
self.draw_point_rgba = np.array(draw_point_rgba, dtype=np.float32)
def _clear_zmp_visualization(self, sim_data: SimData):
sim_data.visual.zmp_world_pos = None
sim_data.visual.zmp_draw_enabled = self.draw_point
sim_data.visual.zmp_draw_size = self.draw_point_size
sim_data.visual.zmp_draw_height_offset = self.draw_height_offset
sim_data.visual.zmp_draw_rgba = self.draw_point_rgba.copy()
def _set_zmp_visualization(self, sim_data: SimData, zmp_world_pos: np.ndarray):
sim_data.visual.zmp_world_pos = np.array(zmp_world_pos, dtype=np.float32)
sim_data.visual.zmp_draw_enabled = self.draw_point
sim_data.visual.zmp_draw_size = self.draw_point_size
sim_data.visual.zmp_draw_height_offset = self.draw_height_offset
sim_data.visual.zmp_draw_rgba = self.draw_point_rgba.copy()
def _log_invalid(self, sim_data: SimData, contact_count: int = 0, metric_value: float = 1.0):
self._clear_zmp_visualization(sim_data)
logger.log(float(metric_value), 'stable_metric/zmp_margin', step=sim_data.n_step)
logger.log(0.0, 'stable_metric/zmp_x', step=sim_data.n_step)
logger.log(0.0, 'stable_metric/zmp_y', step=sim_data.n_step)
logger.log(0.0, 'stable_metric/zmp_norm', step=sim_data.n_step)
logger.log(0.0, 'stable_metric/zmp_d_norm', step=sim_data.n_step)
logger.log(0.0, 'stable_metric/total_force_z', step=sim_data.n_step)
logger.log(0.0, 'stable_metric/support_polygon_area', step=sim_data.n_step)
logger.log(float(contact_count), 'stable_metric/support_contact_count', step=sim_data.n_step)
logger.log(0.0, 'stable_metric/fzmp', step=sim_data.n_step)
logger.log(0.0, 'stable_metric/zmp_valid', step=sim_data.n_step)
def __call__(self, sim_data: SimData, goal_data: GoalData) -> float:
dynamics = sim_data.dynamics
if dynamics is None:
raise RuntimeError("ZMP metric requires sim_data.dynamics, but got None.")
rigid_bodies = dynamics.rigid_bodies
if rigid_bodies is None or rigid_bodies.mass.shape[0] == 0:
raise RuntimeError("ZMP metric requires non-empty dynamics.rigid_bodies.")
contact_mask = dynamics.contacts.distances <= self.contact_threshold
support_contacts = dynamics.contacts.positions[contact_mask]
if support_contacts.shape[0] == 0:
self._log_invalid(sim_data)
return 1.0
support_center = np.mean(support_contacts, axis=0)
rel_com_pos = rigid_bodies.com_pos - support_center[None, :]
gravity = np.asarray(dynamics.gravity, dtype=np.float32)
body_forces = rigid_bodies.mass[:, None] * (gravity[None, :] - rigid_bodies.com_lin_acc)
total_force = np.sum(body_forces, axis=0)
total_force_z = float(total_force[2])
if abs(total_force_z) < self.force_threshold:
self._log_invalid(sim_data, contact_count=support_contacts.shape[0])
return 1.0
inertia_alpha = np.einsum('nij,nj->ni', rigid_bodies.inertia_world, rigid_bodies.ang_acc)
inertia_omega = np.einsum('nij,nj->ni', rigid_bodies.inertia_world, rigid_bodies.ang_vel)
gyro = np.cross(rigid_bodies.ang_vel, inertia_omega)
body_moments = np.cross(rel_com_pos, body_forces) - (inertia_alpha + gyro)
total_moment = np.sum(body_moments, axis=0)
zmp_xy = np.array([
-total_moment[1] / total_force_z,
total_moment[0] / total_force_z
], dtype=np.float32)
zmp_norm = float(np.linalg.norm(zmp_xy))
if dynamics.default_diagonal_foot_distance is None:
raise RuntimeError("ZMP metric requires dynamics.default_diagonal_foot_distance, but got None.")
d_norm = float(dynamics.default_diagonal_foot_distance)
if d_norm < self.D_NORM_MIN:
logger.warning(
f"ZMP metric got too small default_diagonal_foot_distance={d_norm:.8f} "
f"(threshold={self.D_NORM_MIN:.1e}), returning 0.0."
)
self._log_invalid(
sim_data,
contact_count=support_contacts.shape[0],
metric_value=0.0,
)
return 0.0
metric_value = max(0.0, 1.0 - zmp_norm / d_norm)
zmp_world_pos = np.array([
support_center[0] + zmp_xy[0],
support_center[1] + zmp_xy[1],
support_center[2],
], dtype=np.float32)
self._set_zmp_visualization(sim_data, zmp_world_pos)
support_polygon = _convex_hull_2d(support_contacts[:, :2] - support_center[None, :2])
support_area = _polygon_area_2d(support_polygon)
is_inside_support = _point_in_support_region(zmp_xy, support_polygon)
fzmp = 0.0 if is_inside_support else 1.0
logger.log(metric_value, 'stable_metric/zmp_margin', step=sim_data.n_step)
logger.log(float(zmp_xy[0]), 'stable_metric/zmp_x', step=sim_data.n_step)
logger.log(float(zmp_xy[1]), 'stable_metric/zmp_y', step=sim_data.n_step)
logger.log(zmp_norm, 'stable_metric/zmp_norm', step=sim_data.n_step)
logger.log(d_norm, 'stable_metric/zmp_d_norm', step=sim_data.n_step)
logger.log(total_force_z, 'stable_metric/total_force_z', step=sim_data.n_step)
logger.log(float(support_contacts.shape[0]), 'stable_metric/support_contact_count', step=sim_data.n_step)
logger.log(float(support_area), 'stable_metric/support_polygon_area', step=sim_data.n_step)
logger.log(float(fzmp), 'stable_metric/fzmp', step=sim_data.n_step)
logger.log(1.0, 'stable_metric/zmp_valid', step=sim_data.n_step)
return metric_value

3
robogauge/tasks/pipeline/base_pipeline.py
View File

@@ -107,7 +107,8 @@ class BasePipeline:
if self.sim_cfg.domain_rand.action_delay and i == actions_start_decimation:
self.sim.setup_action(action, p_gains, d_gains, control_type)
sim_data = self.sim.step()
self.gauge.update_metrics(sim_data, goal_data)
if self.gauge.update_metrics(sim_data, goal_data):
self.sim.update_visuals(sim_data)
if not self.first_reset and self.gauge.is_reset(sim_data):
sim_data = self.reset_sim_and_robot(sim_data)
except Exception as e:

2
robogauge/tasks/robots/go2/go2_config.py
View File

@@ -48,7 +48,7 @@ class Go2Config(RobotConfig):
cmd = [2.0, 2.0, 0.25]
class commands(RobotConfig.commands):
lin_vel_x = [-2.0, 2.0] # min max [m/s]
lin_vel_x = [-1.0, 1.0] # min max [m/s]
lin_vel_y = [-1.0, 1.0] # min max [m/s]
lin_vel_z = None # min max [m/s]
ang_vel_roll = None # min max [rad/s]

320
robogauge/tasks/simulator/mujoco_simulator.py
View File

@@ -24,7 +24,8 @@ from robogauge.utils.math_utils import get_projected_gravity, quat_rotate_invers
from robogauge.tasks.simulator.mujoco_config import MujocoConfig
from robogauge.tasks.simulator.sim_data import (
SimData,
RobotProprioception, JointState, BaseState, IMUState
RobotProprioception, JointState, BaseState, IMUState,
RigidBodyDynamics, GroundContactState, DynamicsState, VisualState
)
from robogauge.tasks.gauge.goal_data import VelocityGoal
@@ -48,6 +49,25 @@ class MujocoSimulator:
self.target_pos = None
self.target_velocity: Optional[VelocityGoal] = None
self.penetration_reset_count = 0
self.robot_name_prefix = None
self.robot_body_ids = None
self.robot_body_id_set = set()
self.dynamic_body_ids = None
self.dynamic_body_names = []
self.dynamic_body_mass = None
self.dynamic_body_inertia_local = None
self.dynamic_body_iquat = None
self.foot_body_ids = None
self.default_diagonal_foot_distance = 0.0
self.prev_body_com_pos = None
self.prev_body_rot = None
self.prev_body_com_lin_vel = None
self.prev_body_ang_vel = None
self.zmp_world_pos = None
self.zmp_draw_enabled = False
self.zmp_draw_size = 0.03
self.zmp_draw_height_offset = 0.02
self.zmp_draw_rgba = np.array([1.0, 0.85, 0.1, 1.0], dtype=np.float32)
def load(
self,
@@ -114,6 +134,7 @@ class MujocoSimulator:
self.mj_data.qpos[3] = 0.0
self.mj_data.qpos[6] = 1.0
self.mj_data.qpos[7:] = default_dof_pos
self.robot_name_prefix = f'{robot_mjcf.model}/'
# Domain randomization: base mass
base_body_name = f'{robot_mjcf.model}/base_link'
@@ -184,6 +205,8 @@ class MujocoSimulator:
self.penetration_reset_count = 0
self.load_dof_limits()
self.preload_sensors()
self.preload_dynamics_data()
self.reset_dynamics_cache()
# Robot controller placeholders
self.action = None
@@ -259,12 +282,16 @@ class MujocoSimulator:
logger.log(value=proprio.base.lin_vel[1], tag="sim/base_lin_vel_y", step=self.n_step)
if self.n_step == 0:
self.debug_print_proprio_shapes()
dynamics = self.get_dynamics_data()
visual = VisualState()
sim_data = SimData(
n_step=self.n_step,
sim_dt=self.sim_dt,
sim_time=self.sim_time,
proprio=proprio
proprio=proprio,
dynamics=dynamics,
visual=visual,
)
# input("DEBUG")
@@ -288,6 +315,18 @@ class MujocoSimulator:
rgba=[1, 0, 0, 1]
)
def add_zmp_sphere(geom_elem):
zmp_pos = np.array(self.zmp_world_pos, dtype=np.float32).copy()
zmp_pos[2] += self.zmp_draw_height_offset
mujoco.mjv_initGeom(
geom_elem,
type=mujoco.mjtGeom.mjGEOM_SPHERE,
size=[self.zmp_draw_size, 0, 0],
pos=zmp_pos,
mat=np.eye(3).flatten(),
rgba=np.array(self.zmp_draw_rgba, dtype=np.float32)
)
def add_thick_arrow(geom_elem, pos, vec, rgba, scale=0.7):
vel_norm = np.linalg.norm(vec)
display_norm = min(vel_norm * scale, 1.0)
@@ -330,6 +369,14 @@ class MujocoSimulator:
handle.scene.ngeom += 1
add_target_sphere(handle.scene.geoms[self.renderer.scene.ngeom - 1])
if self.zmp_draw_enabled and self.zmp_world_pos is not None:
if ctype == 'viewer':
add_zmp_sphere(handle.user_scn.geoms[viewer_geom_idx])
viewer_geom_idx += 1
else:
handle.scene.ngeom += 1
add_zmp_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]
@@ -413,8 +460,10 @@ class MujocoSimulator:
self.mj_data.qpos[6] = 1.0
self.mj_data.qpos[7:] = self.default_dof_pos
mujoco.mj_forward(self.mj_model, self.mj_data)
self.reset_dynamics_cache()
self.action = None
self.zmp_world_pos = None
if self.viewer is not None:
self.viewer.sync()
@@ -544,6 +593,273 @@ class MujocoSimulator:
data_list.append(self.mj_data.sensordata[adr:adr+dim])
return np.concatenate(data_list)
def update_visuals(self, sim_data: SimData):
visual = sim_data.visual
if visual is None:
self.zmp_world_pos = None
self.zmp_draw_enabled = False
return
self.zmp_world_pos = visual.zmp_world_pos
self.zmp_draw_enabled = bool(visual.zmp_draw_enabled)
self.zmp_draw_size = float(visual.zmp_draw_size)
self.zmp_draw_height_offset = float(visual.zmp_draw_height_offset)
if visual.zmp_draw_rgba is not None:
self.zmp_draw_rgba = np.array(visual.zmp_draw_rgba, dtype=np.float32)
@staticmethod
def _quat_to_rotmat(quat: np.ndarray) -> np.ndarray:
quat = np.asarray(quat, dtype=np.float64)
if quat.ndim == 1:
quat = quat[None, :]
quat_norm = np.linalg.norm(quat, axis=1, keepdims=True)
quat_norm = np.maximum(quat_norm, 1e-12)
quat = quat / quat_norm
w, x, y, z = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3]
rot = np.empty((quat.shape[0], 3, 3), dtype=np.float32)
rot[:, 0, 0] = 1.0 - 2.0 * (y * y + z * z)
rot[:, 0, 1] = 2.0 * (x * y - z * w)
rot[:, 0, 2] = 2.0 * (x * z + y * w)
rot[:, 1, 0] = 2.0 * (x * y + z * w)
rot[:, 1, 1] = 1.0 - 2.0 * (x * x + z * z)
rot[:, 1, 2] = 2.0 * (y * z - x * w)
rot[:, 2, 0] = 2.0 * (x * z - y * w)
rot[:, 2, 1] = 2.0 * (y * z + x * w)
rot[:, 2, 2] = 1.0 - 2.0 * (x * x + y * y)
return rot
@staticmethod
def _rotation_delta_to_angular_velocity(
prev_rot: np.ndarray,
cur_rot: np.ndarray,
dt: float
) -> np.ndarray:
ang_vel = np.zeros((cur_rot.shape[0], 3), dtype=np.float32)
if prev_rot is None or dt <= 0.0:
return ang_vel
rot_delta = cur_rot @ np.transpose(prev_rot, (0, 2, 1))
traces = np.trace(rot_delta, axis1=1, axis2=2)
cos_theta = np.clip((traces - 1.0) * 0.5, -1.0, 1.0)
theta = np.arccos(cos_theta)
skew = np.stack([
rot_delta[:, 2, 1] - rot_delta[:, 1, 2],
rot_delta[:, 0, 2] - rot_delta[:, 2, 0],
rot_delta[:, 1, 0] - rot_delta[:, 0, 1],
], axis=1)
small_mask = theta < 1e-6
if np.any(small_mask):
ang_vel[small_mask] = 0.5 * skew[small_mask] / dt
normal_mask = ~small_mask
if np.any(normal_mask):
idx = np.where(normal_mask)[0]
sin_theta = np.sin(theta[idx])
singular_mask = np.abs(sin_theta) < 1e-6
if np.any(singular_mask):
singular_idx = idx[singular_mask]
ang_vel[singular_idx] = 0.5 * skew[singular_idx] / dt
regular_idx = idx[~singular_mask]
if regular_idx.size > 0:
axis = skew[regular_idx] / (2.0 * sin_theta[~singular_mask, None])
ang_vel[regular_idx] = axis * theta[regular_idx, None] / dt
return ang_vel
@staticmethod
def _max_pairwise_xy_distance(positions: np.ndarray) -> float:
if positions.shape[0] < 2:
return 0.0
xy = positions[:, :2]
diffs = xy[:, None, :] - xy[None, :, :]
dist = np.linalg.norm(diffs, axis=-1)
return float(np.max(dist))
def _get_body_name(self, body_id: int) -> str:
if body_id == 0:
return "worldbody"
name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, body_id)
return name if name is not None else f"body_{body_id}"
def _get_geom_name(self, geom_id: int) -> str:
name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, geom_id)
return name if name is not None else f"geom_{geom_id}"
def reset_dynamics_cache(self):
self.prev_body_com_pos = None
self.prev_body_rot = None
self.prev_body_com_lin_vel = None
self.prev_body_ang_vel = None
def preload_dynamics_data(self):
self.robot_body_ids = []
robot_body_names = []
for body_id in range(self.mj_model.nbody):
body_name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, body_id)
if body_name is not None and body_name.startswith(self.robot_name_prefix):
self.robot_body_ids.append(body_id)
robot_body_names.append(body_name)
self.robot_body_ids = np.array(self.robot_body_ids, dtype=np.int32)
self.robot_body_id_set = set(int(body_id) for body_id in self.robot_body_ids.tolist())
if self.robot_body_ids.size == 0:
logger.warning("No robot bodies found for dynamics preprocessing.")
self.dynamic_body_ids = np.array([], dtype=np.int32)
self.dynamic_body_names = []
self.dynamic_body_mass = np.zeros((0,), dtype=np.float32)
self.dynamic_body_inertia_local = np.zeros((0, 3), dtype=np.float32)
self.dynamic_body_iquat = np.zeros((0, 4), dtype=np.float32)
self.foot_body_ids = np.array([], dtype=np.int32)
self.default_diagonal_foot_distance = 0.0
return
robot_body_mass = np.array(self.mj_model.body_mass[self.robot_body_ids], dtype=np.float32)
dynamic_mask = robot_body_mass > 0.0
self.dynamic_body_ids = self.robot_body_ids[dynamic_mask]
self.dynamic_body_names = [
robot_body_names[idx]
for idx, is_dynamic in enumerate(dynamic_mask.tolist())
if is_dynamic
]
self.dynamic_body_mass = np.array(self.mj_model.body_mass[self.dynamic_body_ids], dtype=np.float32)
self.dynamic_body_inertia_local = np.array(self.mj_model.body_inertia[self.dynamic_body_ids], dtype=np.float32)
if hasattr(self.mj_model, 'body_iquat'):
self.dynamic_body_iquat = np.array(self.mj_model.body_iquat[self.dynamic_body_ids], dtype=np.float32)
else:
self.dynamic_body_iquat = np.tile(np.array([[1.0, 0.0, 0.0, 0.0]], dtype=np.float32), (len(self.dynamic_body_ids), 1))
self.foot_body_ids = np.array([
body_id for body_id, body_name in zip(self.robot_body_ids.tolist(), robot_body_names)
if 'foot' in body_name.rsplit('/', 1)[-1].lower()
], dtype=np.int32)
self.default_diagonal_foot_distance = self.compute_default_diagonal_foot_distance()
logger.info(
f"Dynamics preprocessing: robot_bodies={len(self.robot_body_ids)}, "
f"dynamic_bodies={len(self.dynamic_body_ids)}, foot_bodies={len(self.foot_body_ids)}, "
f"default_diagonal_foot_distance={self.default_diagonal_foot_distance:.4f}"
)
def compute_default_diagonal_foot_distance(self) -> float:
if self.foot_body_ids is not None and self.foot_body_ids.size >= 2:
if hasattr(self.mj_data, 'xpos'):
foot_pos = np.array(self.mj_data.xpos[self.foot_body_ids], dtype=np.float32)
else:
foot_pos = np.array(self.mj_data.xipos[self.foot_body_ids], dtype=np.float32)
return self._max_pairwise_xy_distance(foot_pos)
contacts = self.get_ground_contact_state()
if contacts.positions.shape[0] >= 2:
return self._max_pairwise_xy_distance(contacts.positions)
return 0.0
def get_ground_contact_state(self) -> GroundContactState:
positions = []
distances = []
robot_geom_names = []
other_geom_names = []
robot_body_names = []
other_body_names = []
for i in range(self.mj_data.ncon):
contact = self.mj_data.contact[i]
geom1 = int(contact.geom1)
geom2 = int(contact.geom2)
body1 = int(self.mj_model.geom_bodyid[geom1])
body2 = int(self.mj_model.geom_bodyid[geom2])
is_robot1 = body1 in self.robot_body_id_set
is_robot2 = body2 in self.robot_body_id_set
if is_robot1 == is_robot2:
continue
robot_geom = geom1 if is_robot1 else geom2
other_geom = geom2 if is_robot1 else geom1
robot_body = body1 if is_robot1 else body2
other_body = body2 if is_robot1 else body1
positions.append(np.array(contact.pos, dtype=np.float32))
distances.append(float(contact.dist))
robot_geom_names.append(self._get_geom_name(robot_geom))
other_geom_names.append(self._get_geom_name(other_geom))
robot_body_names.append(self._get_body_name(robot_body))
other_body_names.append(self._get_body_name(other_body))
positions = np.array(positions, dtype=np.float32).reshape(-1, 3)
distances = np.array(distances, dtype=np.float32)
return GroundContactState(
positions=positions,
distances=distances,
robot_geom_names=robot_geom_names,
other_geom_names=other_geom_names,
robot_body_names=robot_body_names,
other_body_names=other_body_names,
)
def get_dynamics_data(self) -> DynamicsState:
if self.dynamic_body_ids is None or self.dynamic_body_ids.size == 0:
empty_body = RigidBodyDynamics(
names=[],
mass=np.zeros((0,), dtype=np.float32),
com_pos=np.zeros((0, 3), dtype=np.float32),
com_lin_vel=np.zeros((0, 3), dtype=np.float32),
com_lin_acc=np.zeros((0, 3), dtype=np.float32),
ang_vel=np.zeros((0, 3), dtype=np.float32),
ang_acc=np.zeros((0, 3), dtype=np.float32),
inertia_world=np.zeros((0, 3, 3), dtype=np.float32),
)
return DynamicsState(
gravity=np.array(self.mj_model.opt.gravity, dtype=np.float32),
rigid_bodies=empty_body,
contacts=self.get_ground_contact_state(),
default_diagonal_foot_distance=float(self.default_diagonal_foot_distance),
)
com_pos = np.array(self.mj_data.xipos[self.dynamic_body_ids], dtype=np.float32)
body_rot = np.array(self.mj_data.xmat[self.dynamic_body_ids], dtype=np.float32).reshape(-1, 3, 3)
com_lin_vel = np.zeros_like(com_pos)
com_lin_acc = np.zeros_like(com_pos)
ang_vel = np.zeros_like(com_pos)
ang_acc = np.zeros_like(com_pos)
if self.prev_body_com_pos is not None:
com_lin_vel = (com_pos - self.prev_body_com_pos) / self.sim_dt
if self.prev_body_com_lin_vel is not None:
com_lin_acc = (com_lin_vel - self.prev_body_com_lin_vel) / self.sim_dt
if self.prev_body_rot is not None:
ang_vel = self._rotation_delta_to_angular_velocity(self.prev_body_rot, body_rot, self.sim_dt)
if self.prev_body_ang_vel is not None:
ang_acc = (ang_vel - self.prev_body_ang_vel) / self.sim_dt
if hasattr(self.mj_data, 'ximat'):
inertial_rot = np.array(self.mj_data.ximat[self.dynamic_body_ids], dtype=np.float32).reshape(-1, 3, 3)
else:
inertial_rot = body_rot @ self._quat_to_rotmat(self.dynamic_body_iquat)
inertia_local = np.zeros((self.dynamic_body_ids.size, 3, 3), dtype=np.float32)
inertia_local[:, 0, 0] = self.dynamic_body_inertia_local[:, 0]
inertia_local[:, 1, 1] = self.dynamic_body_inertia_local[:, 1]
inertia_local[:, 2, 2] = self.dynamic_body_inertia_local[:, 2]
inertia_world = inertial_rot @ inertia_local @ np.transpose(inertial_rot, (0, 2, 1))
self.prev_body_com_pos = com_pos.copy()
self.prev_body_rot = body_rot.copy()
self.prev_body_com_lin_vel = com_lin_vel.copy()
self.prev_body_ang_vel = ang_vel.copy()
return DynamicsState(
gravity=np.array(self.mj_model.opt.gravity, dtype=np.float32),
rigid_bodies=RigidBodyDynamics(
names=self.dynamic_body_names,
mass=self.dynamic_body_mass.copy(),
com_pos=com_pos,
com_lin_vel=com_lin_vel,
com_lin_acc=com_lin_acc,
ang_vel=ang_vel,
ang_acc=ang_acc,
inertia_world=inertia_world,
),
contacts=self.get_ground_contact_state(),
default_diagonal_foot_distance=float(self.default_diagonal_foot_distance),
)
def debug_print_proprio_shapes(self):
"""Log shapes (or lengths) of each numpy vector inside a RobotProprioception.

37
robogauge/tasks/simulator/sim_data.py
View File

@@ -30,9 +30,46 @@ class RobotProprioception:
base: BaseState
imu: IMUState
@dataclass
class RigidBodyDynamics:
names: list # list of rigid body names
mass: np.ndarray # [kg] shape (n_body,)
com_pos: np.ndarray # [m] world frame, shape (n_body, 3)
com_lin_vel: np.ndarray # [m/s] world frame, shape (n_body, 3)
com_lin_acc: np.ndarray # [m/s^2] world frame, shape (n_body, 3)
ang_vel: np.ndarray # [rad/s] world frame, shape (n_body, 3)
ang_acc: np.ndarray # [rad/s^2] world frame, shape (n_body, 3)
inertia_world: np.ndarray # [kg*m^2] world frame, shape (n_body, 3, 3)
@dataclass
class GroundContactState:
positions: np.ndarray # [m] world frame, shape (n_contact, 3)
distances: np.ndarray # [m] shape (n_contact,)
robot_geom_names: list # list of robot geom names
other_geom_names: list # list of non-robot geom names
robot_body_names: list # list of robot body names
other_body_names: list # list of non-robot body names
@dataclass
class DynamicsState:
gravity: np.ndarray
rigid_bodies: RigidBodyDynamics
contacts: GroundContactState
default_diagonal_foot_distance: float
@dataclass
class VisualState:
zmp_world_pos: np.ndarray = None
zmp_draw_enabled: bool = False
zmp_draw_size: float = 0.03
zmp_draw_height_offset: float = 0.02
zmp_draw_rgba: np.ndarray = None
@dataclass
class SimData:
n_step: int
sim_dt: float
sim_time: float
proprio: RobotProprioception
dynamics: DynamicsState
visual: VisualState

94
scripts/run_models.bash Executable file
View File

@@ -0,0 +1,94 @@
#!/bin/bash
source /data/user/wutianyang/Programs/miniforge3/bin/activate go2-gym
python robogauge/scripts/run.py \
--task-name go2_moe \
--model-path /data/user/wutianyang/Coding/go2_rl_gym/mytest_merge_data/rem_cts/go2_moe_cts_expert_goal_137000_0.6745/policies/policy.pt \
--experiment-name go2_moe_cts_expert_goal_137000_0.6745 \
--stress-benchmark \
--stress-terrain-names flat slope_fd slope_bd stairs_fd stairs_bd wave obstacle \
--num-processes 70 \
--seeds 0 1 2 \
--search-seeds 0 1 2 3 4 \
--frictions 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 \
--compress-logs \
--headless
python robogauge/scripts/run.py \
--task-name go2_moe \
--model-path /data/user/wutianyang/Coding/go2_rl_gym/mytest_merge_data/cts_vanilla2/go2_cts_vanilla2_103.5k_0.5819/policies/policy.pt \
--experiment-name go2_cts_vanilla2_103.5k_0.5819 \
--stress-benchmark \
--stress-terrain-names flat slope_fd slope_bd stairs_fd stairs_bd wave obstacle \
--num-processes 70 \
--seeds 0 1 2 \
--search-seeds 0 1 2 3 4 \
--frictions 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 \
--compress-logs \
--headless
python robogauge/scripts/run.py \
--task-name go2 \
--model-path /data/user/wutianyang/Coding/go2_rl_gym/mytest_merge_data/him/go2_him_21k_0.5266/policies/policy.pt \
--experiment-name go2_him_21k_0.5266 \
--stress-benchmark \
--stress-terrain-names flat slope_fd slope_bd stairs_fd stairs_bd wave obstacle \
--num-processes 70 \
--seeds 0 1 2 \
--search-seeds 0 1 2 3 4 \
--frictions 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 \
--compress-logs \
--headless
python robogauge/scripts/run.py \
--task-name go2 \
--model-path /data/user/wutianyang/Coding/go2_rl_gym/mytest_merge_data/dwaq/go2_dwaq_119.5k_0.4691/policies/policy.pt \
--experiment-name go2_dwaq_119.5k_0.4691 \
--stress-benchmark \
--stress-terrain-names flat slope_fd slope_bd stairs_fd stairs_bd wave obstacle \
--num-processes 70 \
--seeds 0 1 2 \
--search-seeds 0 1 2 3 4 \
--frictions 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 \
--compress-logs \
--headless
python robogauge/scripts/run.py \
--task-name go2_moe \
--model-path /data/user/wutianyang/Coding/go2_rl_gym/mytest_merge_data/moe_cts/go2_moe_cts_79k_0.6637/policies/policy.pt \
--experiment-name go2_moe_cts_79k_0.6637 \
--stress-benchmark \
--stress-terrain-names flat slope_fd slope_bd stairs_fd stairs_bd wave obstacle \
--num-processes 70 \
--seeds 0 1 2 \
--search-seeds 0 1 2 3 4 \
--frictions 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 \
--compress-logs \
--headless
python robogauge/scripts/run.py \
--task-name go2_moe \
--model-path /data/user/wutianyang/Coding/go2_rl_gym/mytest_merge_data/ac_moe_cts/go2_ac_moe_cts_115k_0.6589/policies/policy.pt \
--experiment-name go2_ac_moe_cts_115k_0.6589 \
--stress-benchmark \
--stress-terrain-names flat slope_fd slope_bd stairs_fd stairs_bd wave obstacle \
--num-processes 70 \
--seeds 0 1 2 \
--search-seeds 0 1 2 3 4 \
--frictions 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 \
--compress-logs \
--headless
python robogauge/scripts/run.py \
--task-name go2_moe \
--model-path /data/user/wutianyang/Coding/go2_rl_gym/mytest_merge_data/mcp_cts/go2_mcp_cts_91k_0.6513/policies/policy.pt \
--experiment-name go2_mcp_cts_91k_0.6513 \
--stress-benchmark \
--stress-terrain-names flat slope_fd slope_bd stairs_fd stairs_bd wave obstacle \
--num-processes 70 \
--seeds 0 1 2 \
--search-seeds 0 1 2 3 4 \
--frictions 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 \
--compress-logs \
--headless

8
setup.py
View File

@@ -1,12 +1,12 @@
from setuptools import setup, find_packages
setup(
name="robogauge", # 包名
version="1.1.3", # 版本号
author="Wu Tianyang", # 你的名字
name="robogauge",
version="1.1.4",
author="Wu Tianyang",
author_email="993660140@qq.com",
description="A generic robot RL model evaluation library based on MuJoCo",
url="https://github.com/wty-yy/robot_gauge", # 如果有仓库地址
url="https://github.com/wty-yy/robot_gauge",
packages=find_packages(),
install_requires=[
"torch", # Refer: https://pytorch.org/get-started/locally/