Files
sim2sim_mujoco_example/mujoco_api_demo.py
2026-04-05 07:37:08 +08:00

231 lines
9.7 KiB
Python

#!/usr/bin/env python3
"""
MuJoCo API 核心用法演示 - 关节信息读取与控制
"""
import mujoco
import numpy as np
# 加载模型
XML_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/xml/go1.xml"
MESH_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/meshes"
with open(XML_PATH, 'r') as f:
xml = f.read()
xml = xml.replace('meshdir="../meshes/"', f'meshdir="{MESH_PATH}"')
xml = xml.replace('<default>\n=', '<default>\n')
model = mujoco.MjModel.from_xml_string(xml)
data = mujoco.MjData(model)
# ============================================================
# 1. 获取模型基本信息
# ============================================================
print("=" * 60)
print("1. 模型基本信息")
print("=" * 60)
print(f" nq (自由度数量): {model.nq}") # 位置自由度
print(f" nv (速度自由度): {model.nv}") # 速度自由度
print(f" nu (控制数量): {model.nu}") # actuator 数量
print(f" njnt (关节数量): {model.njnt}") # 关节数量
# ============================================================
# 2. 读取关节信息
# ============================================================
print("\n" + "=" * 60)
print("2. 关节信息")
print("=" * 60)
# 关节索引 0 是 world 或 freejoint, 从 1 开始是实际关节
print("\n关节列表 (用 mj_id2name):")
for i in range(model.njnt):
name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i)
print(f" [{i:2d}] {name}")
# ============================================================
# 3. 状态向量详解
# ============================================================
print("\n" + "=" * 60)
print("3. 状态向量")
print("=" * 60)
print(f"""
MuJoCo 状态向量布局 (Go1 是自由关节机器人):
qpos (位置, {model.nq}维):
qpos[0:3] = 全局位置 (x, y, z)
qpos[3:7] = 全局四元数 (qx, qy, qz, qw)
qpos[7:19] = 12 个关节位置 (hip, thigh, calf)
qpos[19:] = 额外数据
qvel (速度, {model.nv}维):
qvel[0:3] = 全局线速度
qvel[3:6] = 全局角速度
qvel[6:18] = 12 个关节速度
""")
print("当前状态:")
print(f" qpos = {data.qpos}")
print(f" qvel = {data.qvel}")
# ============================================================
# 4. 通过关节名读取特定关节数据
# ============================================================
print("\n" + "=" * 60)
print("4. 通过关节名读取数据")
print("=" * 60)
# 方法: mj_name2id 获取关节 ID, 然后用 jnt_qposadr/jnt_dofadr 找到索引
def get_joint_state(model, data, joint_name):
"""获取指定关节的位置和速度
注意: MuJoCo 中:
- qpos 位置索引通过 jnt_qposadr 获取
- qvel 速度索引通过 jnt_dofadr 获取 (不是 jnt_qveladr!)
"""
joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, joint_name)
if joint_id < 0:
return None, None
qposadr = model.jnt_qposadr[joint_id]
dofadr = model.jnt_dofadr[joint_id] # 注意: 是 dofadr 不是 qveladr
# 读取位置 (每个关节 1 个 qpos 值, freejoint 是 7 个)
nq = model.jnt_type[joint_id] == mujoco.mjtJoint.mjJNT_FREE and 7 or 1
pos = data.qpos[qposadr:qposadr+nq]
# 读取速度 (每个关节 1 个 dof, freejoint 没有速度直接索引)
vel = data.qvel[dofadr] if dofadr >= 0 else 0.0
return pos, vel
# 示例: 读取前腿关节
print("\n读取前右腿关节:")
for name in ["FR_hip_joint", "FR_thigh_joint", "FR_calf_joint"]:
pos, vel = get_joint_state(model, data, name)
pos_str = f"{pos:.4f}" if np.isscalar(pos) else str(pos)
print(f" {name:20s}: pos={pos_str}, vel={vel:.4f}")
# ============================================================
# 5. 控制输出
# ============================================================
print("\n" + "=" * 60)
print("5. 控制输出 (ctrl)")
print("=" * 60)
print(f"\n控制向量长度: {model.nu}")
print(f"控制范围 (从 XML 读取):\n{model.actuator_ctrlrange}")
# 遍历所有 actuator
print("\nActuator 列表:")
for i in range(model.nu):
name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i)
# 获取 actuator 关联的 joint
trn_joint = model.actuator_trnid[i, 0]
joint_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, trn_joint) if trn_joint >= 0 else "none"
ctrl_range = model.actuator_ctrlrange[i]
print(f" [{i:2d}] {name:15s} -> {joint_name:15s} range=[{ctrl_range[0]:.2f}, {ctrl_range[1]:.2f}]")
# ============================================================
# 6. 控制输出方法
# ============================================================
print("\n" + "=" * 60)
print("6. 控制输出方法")
print("=" * 60)
print("""
# 方法1: 直接设置所有控制
data.ctrl[:] = 0.0
# 方法2: 设置单个控制
data.ctrl[0] = 1.0
# 方法3: 用索引设置 (与 actuator 顺序对应)
data.ctrl[np.arange(12)] = target_positions
# 方法4: 先计算再设置
target_torque = kp * (desired - actual) + kd * (desired_vel - actual_vel)
data.ctrl[:] = target_torque
""")
# 示例: 设置目标关节位置
desired_joint_pos = np.zeros(12)
for i in range(12):
desired_joint_pos[i] = np.random.uniform(-0.5, 0.5)
print(f"\n示例: 设置目标关节位置: {desired_joint_pos}")
data.ctrl[:] = desired_joint_pos
print(f"已写入 ctrl: {data.ctrl}")
# ============================================================
# 7. 仿真循环
# ============================================================
print("\n" + "=" * 60)
print("7. 仿真循环示例")
print("=" * 60)
print("""
# 基本仿真循环:
for _ in range(100):
# 1. 读取当前状态
current_pos = data.qpos[7:19] # 关节位置
current_vel = data.qvel[6:18] # 关节速度
# 2. 计算控制量 (这里用简单的 PD 控制作为示例)
desired_pos = np.array([...]) # 目标位置
error = desired_pos - current_pos
data.ctrl[:] = kp * error - kd * current_vel
# 3. 执行一步仿真
mujoco.mj_step(model, data)
# 4. (可选) 渲染
# renderer.update_scene(data)
# img = renderer.render()
""")
# 实际运行 50 步
print("\n运行 50 步仿真 (零控制):")
data.ctrl[:] = 0.0
for step in range(50):
mujoco.mj_step(model, data)
if step % 10 == 0:
print(f" Step {step:3d}: trunk_z={data.qpos[2]:.4f}m, "
f"FR_hip={data.qpos[7]:.4f}, FR_thigh={data.qpos[8]:.4f}")
print("\n机器人倒下是因为零控制 + 重力!")
print("真实 RL 控制需要在每步根据策略计算正确的 ctrl 值。")
# ============================================================
# 8. 重要 API 速查
# ============================================================
print("\n" + "=" * 60)
print("8. 重要 API 速查")
print("=" * 60)
print("""
┌─────────────────────────────────────────────────────────────┐
│ 读取模型信息 │
├─────────────────────────────────────────────────────────────┤
│ model.nq # 自由度数 │
│ model.nu # 控制数 │
│ model.njnt # 关节数 │
│ mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i) │
│ mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name) │
├─────────────────────────────────────────────────────────────┤
│ 读取状态 │
├─────────────────────────────────────────────────────────────┤
│ data.qpos # 位置向量 │
│ data.qvel # 速度向量 │
│ data.qacc # 加速度向量 │
│ data.ctrl # 控制向量 (要设置的) │
├─────────────────────────────────────────────────────────────┤
│ 关节层面 │
├─────────────────────────────────────────────────────────────┤
│ joint_id = mujoco.mj_name2id(model, ..., joint_name) │
│ qpos_adr = model.jnt_qposadr[joint_id] │
│ qvel_adr = model.jnt_qveladr[joint_id] │
│ joint_pos = data.qpos[qpos_adr] │
│ joint_vel = data.qvel[qvel_adr] │
├─────────────────────────────────────────────────────────────┤
│ 仿真 │
├─────────────────────────────────────────────────────────────┤
│ mujoco.mj_step(model, data) # 执行一步 │
│ mujoco.mj_stepN(model, data, N) # 执行 N 步 │
│ mujoco.mj_resetData(model, data) # 重置状态 │
└─────────────────────────────────────────────────────────────┘
""")