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

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)
@@ -329,6 +368,14 @@ class MujocoSimulator:
else:
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]
@@ -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.

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

@@ -5,7 +5,7 @@ from dataclasses import dataclass
class JointState:
pos: np.ndarray # [rad] shape (n_dof,)
vel: np.ndarray # [rad/s] shape (n_dof,)
torque: np.ndarray # [N*m] shape (n_dof,)
torque: np.ndarray # [N*m] shape (n_dof,)
limits: np.ndarray # [rad] shape (n_dof, 2), lower and upper limits
names: list # list of joint names
@@ -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