wtw init
This commit is contained in:
cyylinux
9
.claude/settings.json
Normal file
9
.claude/settings.json
Normal file
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"permissions": {
|
||||
"allow": [
|
||||
"Bash(curl -s \"https://api.github.com/repos/zzzJie-Robot/Go2Arm_Lab/contents/source/Go2Arm_Lab/Go2Arm_Lab/tasks/manager_based/go2arm_lab\")",
|
||||
"Read(//home/8x54zj-m/walk-these-ways/**)",
|
||||
"Bash(python3 -c \"import ast; ast.parse\\(open\\('go1_walk_these_ways_inference.py'\\).read\\(\\)\\)\")"
|
||||
]
|
||||
}
|
||||
}
|
||||
BIN
__pycache__/go1_stairs_sim2sim_inference.cpython-313.pyc
Normal file
BIN
__pycache__/go1_stairs_sim2sim_inference.cpython-313.pyc
Normal file
Binary file not shown.
BIN
__pycache__/go1_stairs_sim2sim_inference.cpython-38.pyc
Normal file
BIN
__pycache__/go1_stairs_sim2sim_inference.cpython-38.pyc
Normal file
Binary file not shown.
BIN
__pycache__/go1_unitree_rl_inference.cpython-310.pyc
Normal file
BIN
__pycache__/go1_unitree_rl_inference.cpython-310.pyc
Normal file
Binary file not shown.
BIN
__pycache__/go1_walk_these_ways_inference.cpython-310.pyc
Normal file
BIN
__pycache__/go1_walk_these_ways_inference.cpython-310.pyc
Normal file
Binary file not shown.
BIN
__pycache__/go1_walk_these_ways_inference.cpython-38.pyc
Normal file
BIN
__pycache__/go1_walk_these_ways_inference.cpython-38.pyc
Normal file
Binary file not shown.
@@ -9,7 +9,7 @@
|
||||
<light castshadow="false" diffuse="1 1 1"/>
|
||||
<motor ctrlrange="-33.5 33.5" ctrllimited="true"/>
|
||||
<camera fovy="60"/>
|
||||
<joint damping="0.01" armature="0.01" frictionloss="0.2" />
|
||||
<joint damping="1.0" armature="0.01" frictionloss="0.5" />
|
||||
|
||||
</default>
|
||||
=
|
||||
|
||||
1
data/go1/xml/terrain/plane.xml
Normal file
1
data/go1/xml/terrain/plane.xml
Normal file
@@ -0,0 +1 @@
|
||||
<geom name='floor' type='plane' conaffinity='1' condim='3' contype='1' rgba="0.5 0.9 0.9 0.1" material='plane' pos='0 0 0' size='0 0 1' friction="0.6 0.3 0.3"/>
|
||||
12
data/go1/xml/terrain/rough.xml
Normal file
12
data/go1/xml/terrain/rough.xml
Normal file
@@ -0,0 +1,12 @@
|
||||
<!-- 模拟碎地形 -->
|
||||
<geom name='floor' type='plane' conaffinity='1' condim='3' contype='1' rgba="0.4 0.6 0.3 1" pos='0 0 0' size='0 0 1' friction="0.6 0.3 0.3"/>
|
||||
<geom name='r1' type='box' size='0.15 0.1 0.03' pos='0.5 0.2 0' rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='r2' type='box' size='0.1 0.15 0.05' pos='0.8 -0.1 0' rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='r3' type='box' size='0.12 0.08 0.04' pos='1.1 0.15 0' rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='r4' type='box' size='0.08 0.12 0.06' pos='1.4 -0.2 0' rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='r5' type='box' size='0.1 0.1 0.035' pos='1.7 0.1 0' rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='r6' type='box' size='0.14 0.09 0.045' pos='2.0 -0.05 0' rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='r7' type='box' size='0.09 0.14 0.055' pos='0.3 0.25 0' rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='r8' type='box' size='0.11 0.07 0.025' pos='0.6 -0.15 0' rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='r9' type='box' size='0.07 0.11 0.05' pos='0.9 0.05 0' rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='r10' type='box' size='0.13 0.08 0.04' pos='1.2 -0.1 0' rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
6
data/go1/xml/terrain/slope.xml
Normal file
6
data/go1/xml/terrain/slope.xml
Normal file
@@ -0,0 +1,6 @@
|
||||
<!-- 斜坡地形-->
|
||||
<geom name='start_platform' type='box' size='10.0 10.0 0.1' pos='0 0 -0.06' rgba="0.5 0.6 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='start_fill' type='box' size='10.0 10.0 0.2' pos='0 0 -0.26' rgba="0.4 0.35 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='slope' type='box' size='2.5 0.3 0.1' pos='6.8 0 0.10' quat='0.258819 0 0 0.965926' rgba="0.5 0.6 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='top_platform' type='box' size='1.0 0.3 0.06' pos='8.8 0 0.16' rgba="0.5 0.6 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='slope_fill' type='box' size='2.5 0.3 0.25' pos='6.8 0 -0.15' rgba="0.4 0.35 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
13
data/go1/xml/terrain/stairs.xml
Normal file
13
data/go1/xml/terrain/stairs.xml
Normal file
@@ -0,0 +1,13 @@
|
||||
<!-- 楼梯地形 -->
|
||||
<geom name='step0' type='box' size='0.5 10.4 0.06' pos='0 0 -0.06' rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='step1' type='box' size='0.5 10.4 0.06' pos='0.3 0 -0.06' rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='step2' type='box' size='0.5 10.4 0.06' pos='0.6 0 0.00' rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='step3' type='box' size='0.5 10.4 0.06' pos='0.9 0 0.06' rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='step4' type='box' size='0.5 10.4 0.06' pos='1.2 0 0.12' rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='step5' type='box' size='0.5 10.4 0.06' pos='1.5 0 0.18' rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='step6' type='box' size='0.5 10.4 0.06' pos='1.8 0 0.24' rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='step7' type='box' size='0.5 10.4 0.06' pos='2.1 0 0.30' rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='step8' type='box' size='0.5 10.4 0.06' pos='2.4 0 0.36' rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='step9' type='box' size='0.5 10.4 0.06' pos='2.7 0 0.42' rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='platform' type='box' size='0.8 10.4 0.06' pos='3.3 0 0.45' rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
|
||||
<geom name='fill' type='box' size='2.0 10.4 0.45' pos='1.5 0 -0.45' rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
@@ -1,11 +1,12 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Go1 RL Policy 推理 - Walk-These-Ways 预训练模型 + MuJoCo 可视化
|
||||
Go1 RL Policy 推理demo
|
||||
默认加载的Walk-These-Ways 预训练模型 + MuJoCo 可视化
|
||||
模型: body_latest.jit + adaptation_module_latest.jit (GRU-based policy with history)
|
||||
训练环境: IsaacGym (walk-these-ways)
|
||||
|
||||
使用方法:
|
||||
python3 /home/8x54zj-m/unitree_mujoco/go1_walk_these_ways_inference.py
|
||||
source /opt/mujoco_project/mujoco_env/bin/activate
|
||||
python3 /opt/unitree_mujoco/go1_walk_these_ways_inference.py
|
||||
|
||||
键盘控制:
|
||||
W/S: 前进/后退
|
||||
@@ -21,16 +22,11 @@ import numpy as np
|
||||
import mujoco
|
||||
from mujoco import viewer
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import signal
|
||||
import collections
|
||||
import queue
|
||||
|
||||
# ============================================================
|
||||
# 全局退出标志
|
||||
# ============================================================
|
||||
g_exit_requested = False
|
||||
|
||||
def signal_handler(signum, frame):
|
||||
@@ -40,14 +36,18 @@ def signal_handler(signum, frame):
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
# ============================================================
|
||||
# 配置
|
||||
# 预训练模型 /opt/walk-these-ways/repo/runs/gait-conditioned-agility/pretrain-v0/train/025417.456545/checkpoints
|
||||
# ============================================================
|
||||
MODEL_DIR = "/home/8x54zj-m/walk-these-ways/runs/gait-conditioned-agility/pretrain-v0/train/025417.456545/checkpoints"
|
||||
MODEL_DIR = "/opt/walk-these-ways/repo/runs/gait-conditioned-agility/pretrain-v0/train/025417.456545/checkpoints"
|
||||
BODY_MODEL_PATH = os.path.join(MODEL_DIR, "body_latest.jit")
|
||||
ADAPT_MODEL_PATH = os.path.join(MODEL_DIR, "adaptation_module_latest.jit")
|
||||
|
||||
XML_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/xml/go1.xml"
|
||||
MESH_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/meshes"
|
||||
XML_PATH = "/opt/unitree_mujoco/data/go1/xml/go1.xml"
|
||||
MESH_PATH = "/opt/unitree_mujoco/data/go1/meshes"
|
||||
TERRAIN_PATH = "/opt/unitree_mujoco/data/go1/xml/terrain"
|
||||
|
||||
# 地形配置: "plane", "stairs", "slope", "rough"
|
||||
TERRAIN = "stairs"
|
||||
|
||||
# ============================================================
|
||||
# Walk-These-Ways 模型参数 (来自 parameters.pkl & deploy配置)
|
||||
@@ -94,7 +94,7 @@ DEPLOY_TO_MUJOCO_MAPPING = np.array([3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8])
|
||||
|
||||
# 速度命令范围
|
||||
MAX_LIN_VEL = 1.0 # 最大线速度 m/s
|
||||
MAX_ANG_VEL = 1.0 # 与 deploy 一致
|
||||
MAX_ANG_VEL = 3.0 # 与 deploy 一致
|
||||
|
||||
# PD控制参数
|
||||
KP = 20.0
|
||||
@@ -186,7 +186,7 @@ class KeyboardReader:
|
||||
self._listener.start()
|
||||
self._reader_thread = threading.Thread(target=self._reader_worker, daemon=True)
|
||||
self._reader_thread.start()
|
||||
print("[INFO] 键盘监听已启动 (独立线程)")
|
||||
print("[INFO] 键盘监听已启动")
|
||||
except Exception as e:
|
||||
print(f"[WARN] 无法初始化键盘监听: {e}")
|
||||
|
||||
@@ -211,85 +211,14 @@ class KeyboardReader:
|
||||
# 辅助函数
|
||||
# ============================================================
|
||||
|
||||
def quaternion_to_rotation_matrix(q):
|
||||
"""四元数转旋转矩阵 (MuJoCo 格式: qx, qy, qz, qw)"""
|
||||
qx, qy, qz, qw = q
|
||||
norm = np.sqrt(qx**2 + qy**2 + qz**2 + qw**2)
|
||||
qx, qy, qz, qw = qx/norm, qy/norm, qz/norm, qw/norm
|
||||
return np.array([
|
||||
[1-2*(qy**2+qz**2), 2*(qx*qy-qz*qw), 2*(qx*qz+qy*qw)],
|
||||
[2*(qx*qy+qz*qw), 1-2*(qx**2+qz**2), 2*(qy*qz-qx*qw)],
|
||||
[2*(qx*qz-qy*qw), 2*(qy*qz+qx*qw), 1-2*(qx**2+qy**2)]
|
||||
])
|
||||
|
||||
|
||||
def quat_rotate_inverse(data, v):
|
||||
"""四元数逆旋转 (世界坐标系 -> 躯干坐标系)
|
||||
使用 data.xmat 直接计算，更可靠
|
||||
"""
|
||||
四元数逆旋转 (世界坐标系 -> 躯干坐标系)
|
||||
"""
|
||||
base_rot = data.xmat[1].reshape(3, 3) # body 1 = trunk
|
||||
return base_rot.T @ np.array(v, dtype=np.float64)
|
||||
|
||||
|
||||
def rotation_matrix_from_quat(data):
|
||||
"""从四元数计算旋转矩阵 (躯干坐标系)
|
||||
使用 MuJoCo 内置函数 mju_quat2Mat，保证与 MuJoCo 内部一致
|
||||
返回: 3x3 旋转矩阵 R, 使得 R * world_vec = body_vec
|
||||
"""
|
||||
quat_raw = data.qpos[3:7] # [x, y, z, w]
|
||||
quat_mju = np.array([quat_raw[3], quat_raw[0], quat_raw[1], quat_raw[2]], dtype=np.float64)
|
||||
R_mju = np.zeros(9, dtype=np.float64)
|
||||
mujoco.mju_quat2Mat(R_mju, quat_mju)
|
||||
return R_mju.reshape(3, 3)
|
||||
|
||||
|
||||
def grav_to_arrow(grav):
|
||||
"""将重力投影向量转为箭头字符串 (躯干坐标系视图)
|
||||
grav: [gx, gy, gz] 世界重力在躯干坐标系下的投影
|
||||
躯干坐标系: X=前, Y=左, Z=上
|
||||
返回: (world_arrow, body_arrow)
|
||||
"""
|
||||
# 世界重力: 始终是 (0, 0, -1) = 纯Z轴负方向 = 向下 ↓
|
||||
world_arrow = "↓"
|
||||
|
||||
# 躯干坐标系重力投影
|
||||
gx, gy, gz = grav
|
||||
|
||||
# 判断主体方向 (忽略很小分量)
|
||||
abs_x, abs_y, abs_z = abs(gx), abs(gy), abs(gz)
|
||||
|
||||
# 标准化
|
||||
total = abs_x + abs_y + abs_z
|
||||
nx, ny, nz = gx / total, gy / total, gz / total
|
||||
|
||||
# Z 分量判断上下
|
||||
if nz > 0.1:
|
||||
z_arrow = "↑"
|
||||
elif nz < -0.1:
|
||||
z_arrow = "↓"
|
||||
else:
|
||||
z_arrow = "·"
|
||||
|
||||
# X 分量判断前后
|
||||
if nx > 0.1:
|
||||
x_arrow = "→"
|
||||
elif nx < -0.1:
|
||||
x_arrow = "←"
|
||||
else:
|
||||
x_arrow = "·"
|
||||
|
||||
# Y 分量判断左右 (正Y=机器左=世界右→所以箭头反向)
|
||||
if ny > 0.1:
|
||||
y_arrow = "←" # 机器左倾
|
||||
elif ny < -0.1:
|
||||
y_arrow = "→" # 机器右倾
|
||||
else:
|
||||
y_arrow = "·"
|
||||
|
||||
body_arrow = f"{x_arrow}{y_arrow}{z_arrow}"
|
||||
return world_arrow, body_arrow
|
||||
|
||||
|
||||
def compute_observations_wtw(data, prev_action, last_action, commands, clock_inputs, default_dof_pos_mujoco, obs_scales):
|
||||
"""
|
||||
计算 Walk-These-Ways 策略的观测向量 (70维)
|
||||
@@ -374,37 +303,35 @@ def main():
|
||||
import torch
|
||||
|
||||
# 1. 加载 MuJoCo 模型
|
||||
import re
|
||||
|
||||
os.chdir(MESH_PATH)
|
||||
with open(XML_PATH, 'r') as f:
|
||||
xml_content = f.read()
|
||||
xml_content = xml_content.replace('meshdir="../meshes/"', f'meshdir="{MESH_PATH}"')
|
||||
xml_content = xml_content.replace('<default>\n=', '<default>\n')
|
||||
|
||||
# ---- 添加重力箭头 mocap body (橙色小球) ----
|
||||
# mocap body 会自动在viewer中渲染，无需手动管理geom
|
||||
grav_arrow_body = '''
|
||||
<body name="grav_arrow" pos="0 0 0" mocap="true">
|
||||
<geom type="sphere" size="0.04 0 0" rgba="1.0 0.5 0.0 0.8" contype="0" conaffinity="0" friction="0 0 0"/>
|
||||
</body>
|
||||
'''
|
||||
# 在 </worldbody> 之前插入
|
||||
xml_content = xml_content.replace('</worldbody>', grav_arrow_body + '</worldbody>')
|
||||
# ---- 地形替换: 使用单独的 terrain XML 文件 ----
|
||||
if TERRAIN != "none":
|
||||
terrain_file = os.path.join(TERRAIN_PATH, f"{TERRAIN}.xml")
|
||||
if os.path.exists(terrain_file):
|
||||
with open(terrain_file, 'r') as f:
|
||||
terrain_xml = f.read()
|
||||
# 删除原有的 floor geom
|
||||
xml_content = re.sub(r'<geom[^>]*name=[^>]*floor[^>]*/>', '', xml_content)
|
||||
xml_content = re.sub(r"<geom[^>]*name=[^>]*floor[^>]*/>", '', xml_content)
|
||||
# 在 </worldbody> 之前插入新地形
|
||||
xml_content = xml_content.replace('</worldbody>', terrain_xml + '</worldbody>')
|
||||
print(f"[INFO] 地形已加载: {terrain_file} (TERRAIN={TERRAIN})")
|
||||
else:
|
||||
print(f"[WARN] 地形文件不存在: {terrain_file}, 使用默认地面")
|
||||
|
||||
model = mujoco.MjModel.from_xml_string(xml_content)
|
||||
data = mujoco.MjData(model)
|
||||
|
||||
# 验证 mocap body 添加成功
|
||||
grav_arrow_mocap_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "grav_arrow")
|
||||
if grav_arrow_mocap_id < 0:
|
||||
print("[WARN] grav_arrow body not found in model!")
|
||||
else:
|
||||
print(f"[INFO] grav_arrow body id: {grav_arrow_mocap_id}, nmocap: {model.nmocap}")
|
||||
# mocap_pos[0] 对应第一个 mocap body（我们添加的 grav_arrow）
|
||||
GRAV_MOCAP_IDX = 0
|
||||
if model.nmocap < 1:
|
||||
print("[WARN] No mocap bodies in model! Gravity arrow disabled.")
|
||||
|
||||
print(f"[INFO] MuJoCo Model: {model.nbody} bodies, {model.nq} DoF, {model.nu} actuators")
|
||||
|
||||
print(f"[INFO] 地形: {TERRAIN} ({os.path.join(TERRAIN_PATH, TERRAIN + '.xml')})")
|
||||
print(f"[INFO] Mujoco关节顺序: {[mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i) for i in range(1, 13)]}")
|
||||
print(f"[INFO] Deploy关节顺序: FL, FR, RL, RR (与Mujoco不同)")
|
||||
|
||||
@@ -413,6 +340,8 @@ def main():
|
||||
if floor_id >= 0:
|
||||
model.geom_friction[floor_id] = FLOOR_FRICTION
|
||||
print(f"[INFO] 地面摩擦系数设置为: {FLOOR_FRICTION}")
|
||||
else:
|
||||
print(f"[INFO] 未找到 floor geom，跳过地面摩擦设置（地形可能已自定义摩擦系数）")
|
||||
|
||||
for i in range(model.ngeom):
|
||||
if i != floor_id:
|
||||
@@ -442,62 +371,6 @@ def main():
|
||||
data.qpos[2] = 0.35 # 抬高躯干
|
||||
mujoco.mj_step(model, data)
|
||||
|
||||
# ---- 用 MuJoCo 内置函数验证旋转矩阵 ----
|
||||
quat_mj = data.qpos[3:7].copy() # [x,y,z,w]
|
||||
print(f" [CHECK] MuJoCo qpos[3:7] = {quat_mj}")
|
||||
# MuJoCo's own rotation matrix from quaternion
|
||||
R_mj = np.zeros(9)
|
||||
mujoco.mju_quat2Mat(R_mj, quat_mj)
|
||||
R_mj = R_mj.reshape(3, 3)
|
||||
print(f" [CHECK] MuJoCo mju_quat2Mat R (identity if q=[1,0,0,0]≈id): \n{R_mj.round(4)}")
|
||||
mujoco.mj_step(model, data)
|
||||
|
||||
# 打印初始化后的姿态确认
|
||||
quat_mj = data.qpos[3:7]
|
||||
grav_init = quat_rotate_inverse(data, np.array([0., 0., -1.]))
|
||||
R_init = rotation_matrix_from_quat(data)
|
||||
quat_std = np.array([quat_mj[3], quat_mj[0], quat_mj[1], quat_mj[2]])
|
||||
print(f" [INIT] quat(MuJoCo)={quat_mj} | quat_as_std={quat_std.round(3)} | grav_world→body={grav_init.round(3)} | trunk_z={data.qpos[2]:.3f}")
|
||||
|
||||
# ---- 3D 可视化: 躯干坐标系 vs 世界坐标系 ----
|
||||
try:
|
||||
import matplotlib.pyplot as plt
|
||||
from mpl_toolkits.mplot3d import Axes3D
|
||||
R_init = rotation_matrix_from_quat(data)
|
||||
fig = plt.figure(figsize=(6, 6))
|
||||
ax = fig.add_subplot(111, projection='3d')
|
||||
# 世界坐标系 (黑色)
|
||||
ax.quiver(0, 0, 0, 1.2, 0, 0, color='k', linewidth=1.5, arrow_length_ratio=0.1)
|
||||
ax.quiver(0, 0, 0, 0, 1.2, 0, color='k', linewidth=1.5, arrow_length_ratio=0.1)
|
||||
ax.quiver(0, 0, 0, 0, 0, 1.2, color='k', linewidth=1.5, arrow_length_ratio=0.1)
|
||||
ax.text(1.3, 0, 0, "Xw(前)", fontsize=9)
|
||||
ax.text(0, 1.3, 0, "Yw(左)", fontsize=9)
|
||||
ax.text(0, 0, 1.3, "Zw(上)", fontsize=9)
|
||||
# 躯干坐标系 (彩色)
|
||||
body_x = R_init[:, 0] * 0.8 # R第一列 = body X轴在世界
|
||||
body_y = R_init[:, 1] * 0.8 # R第二列 = body Y轴在世界
|
||||
body_z = R_init[:, 2] * 0.8 # R第三列 = body Z轴在世界
|
||||
ax.quiver(0, 0, 0, *body_x, color='r', linewidth=2, arrow_length_ratio=0.1)
|
||||
ax.quiver(0, 0, 0, *body_y, color='g', linewidth=2, arrow_length_ratio=0.1)
|
||||
ax.quiver(0, 0, 0, *body_z, color='b', linewidth=2, arrow_length_ratio=0.1)
|
||||
ax.text(body_x[0]*1.2, body_x[1]*1.2, body_x[2]*1.2, "Xb(前)", color='r', fontsize=9)
|
||||
ax.text(body_y[0]*1.2, body_y[1]*1.2, body_y[2]*1.2, "Yb(左)", color='g', fontsize=9)
|
||||
ax.text(body_z[0]*1.2, body_z[1]*1.2, body_z[2]*1.2, "Zb(上)", color='b', fontsize=9)
|
||||
# 重力向量
|
||||
grav_arrow = grav_init * 0.6
|
||||
ax.quiver(0, 0, 0, *grav_arrow, color='orange', linewidth=2.5, arrow_length_ratio=0.1)
|
||||
ax.text(grav_arrow[0]*1.2, grav_arrow[1]*1.2, grav_arrow[2]*1.2, f"g={grav_init.round(2)}", color='orange', fontsize=9)
|
||||
ax.set_xlim([-1.5, 1.5]); ax.set_ylim([-1.5, 1.5]); ax.set_zlim([-1.5, 1.5])
|
||||
ax.set_xlabel("X (世界)"); ax.set_ylabel("Y (世界)"); ax.set_zlabel("Z (世界)")
|
||||
ax.set_title("躯干坐标系 (RGB=XYZ轴) vs 世界坐标系 (K) | 橙色=重力投影")
|
||||
plt.tight_layout()
|
||||
plt.savefig("/home/8x54zj-m/unitree_mujoco/body_frame_axes.png", dpi=150)
|
||||
print(f" [INIT] 坐标系可视化已保存: body_frame_axes.png")
|
||||
print(f" [INIT] 旋转矩阵 R (body←world, 列=body轴在world中):\n{R_init.round(3)}")
|
||||
plt.close()
|
||||
except Exception as e:
|
||||
print(f" [WARN] 可视化失败: {e}")
|
||||
|
||||
# 初始化历史 (填充零)
|
||||
for _ in range(NUM_OBS_HISTORY):
|
||||
dummy_obs = np.zeros(NUM_OBS, dtype=np.float32)
|
||||
@@ -704,37 +577,13 @@ def main():
|
||||
|
||||
mujoco.mj_step(model, data)
|
||||
|
||||
# grav_proj: 世界重力 [0,0,-1] 变换到躯干坐标系
|
||||
grav_proj = quat_rotate_inverse(data, np.array([0., 0., -1.]))
|
||||
|
||||
# ---- 重力箭头可视化: 橙色小球始终指向世界下方 ----
|
||||
torso_pos = data.xpos[1] # body 1 = trunk
|
||||
arrow_pos = torso_pos + np.array([0., 0., -1.]) * 0.15
|
||||
|
||||
# ---- 画橙色小球: 用 mocap body 更新位置 ----
|
||||
if grav_arrow_mocap_id >= 0:
|
||||
data.mocap_pos[GRAV_MOCAP_IDX] = arrow_pos.astype(np.float64)
|
||||
|
||||
view.sync()
|
||||
|
||||
step_count += 1
|
||||
if step_count % 100 == 0:
|
||||
trunk_z = data.qpos[2]
|
||||
lin_vel = np.linalg.norm(data.qvel[0:3])
|
||||
quat_mj = data.qpos[3:7]
|
||||
quat = np.concatenate([quat_mj[3:4], quat_mj[0:3]])
|
||||
grav_proj = quat_rotate_inverse(data, np.array([0., 0., -1.]))
|
||||
quat_raw = data.qpos[3:7]
|
||||
quat_mju = np.array([quat_raw[3], quat_raw[0], quat_raw[1], quat_raw[2]], dtype=np.float64)
|
||||
R_mju = np.zeros(9, dtype=np.float64)
|
||||
mujoco.mju_quat2Mat(R_mju, quat_mju)
|
||||
R = R_mju.reshape(3, 3)
|
||||
world_a, body_a = grav_to_arrow(grav_proj)
|
||||
gait_short = {'trot': 'TROT', 'pace': 'PACE', 'bound': 'BOUND', 'pronk': 'PRONK'}
|
||||
# R 的三列 = body X/Y/Z 在世界坐标系中的方向
|
||||
print(f" Step {step_count}: quat={quat_raw.round(3)} grav={grav_proj.round(2)}{body_a} | "
|
||||
f"R=[{R[0,:].round(2)}, {R[1,:].round(2)}, {R[2,:].round(2)}] | "
|
||||
f"trunk_z={trunk_z:.3f}m, vel={lin_vel:.3f}m/s")
|
||||
print(f" Step {step_count}: trunk_z={trunk_z:.3f}m, vel={lin_vel:.3f}m/s")
|
||||
|
||||
keyboard_reader.restore()
|
||||
view.close()
|
||||
@@ -1,618 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Go1 RL Policy 推理 - 使用 ONNX Runtime + MuJoCo 可视化
|
||||
训练模型: /opt/IsaacLab/logs/rsl_rl/unitree_go1_flat/2026-03-24_00-25-36/model_1999.pt
|
||||
|
||||
使用方法:
|
||||
python3 /home/8x54zj-m/unitree_mujoco/go1_rl_inference.py
|
||||
|
||||
运行此脚本需要:
|
||||
pip install mujoco onnxruntime numpy
|
||||
|
||||
键盘控制 (无需 root):
|
||||
W/S: 前进/后退
|
||||
A/D: 左转/右转
|
||||
Q/E: 侧向左移/右移
|
||||
空格: 停止
|
||||
ESC 或 Ctrl+C: 退出查看器
|
||||
"""
|
||||
import numpy as np
|
||||
import mujoco
|
||||
from mujoco import viewer
|
||||
import os
|
||||
import sys
|
||||
import tty
|
||||
import termios
|
||||
import signal
|
||||
import threading
|
||||
import time
|
||||
|
||||
# ============================================================
|
||||
# 全局退出标志
|
||||
# ============================================================
|
||||
g_exit_requested = False
|
||||
|
||||
def signal_handler(signum, frame):
|
||||
"""处理 Ctrl+C (SIGINT) 信号"""
|
||||
global g_exit_requested
|
||||
g_exit_requested = True
|
||||
|
||||
# 注册信号处理器
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
# ============================================================
|
||||
# 配置
|
||||
# ============================================================
|
||||
ONNX_PATH = "/opt/IsaacLab/logs/rsl_rl/unitree_go1_flat/2026-03-24_00-25-36/exported/policy.onnx"
|
||||
XML_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/xml/go1.xml"
|
||||
MESH_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/meshes"
|
||||
MLP_PATH = "/opt/MotrixLab/walk-these-ways/resources/actuator_nets/unitree_go1.pt"
|
||||
|
||||
# 观测/动作维度
|
||||
OBS_DIM = 48
|
||||
ACTION_DIM = 12
|
||||
|
||||
# 关节顺序映射 (policy输出 -> MuJoCo关节)
|
||||
# MuJoCo 关节顺序: [FR_hip, FR_thigh, FR_calf, FL_hip, FL_thigh, FL_calf, RR_hip, RR_thigh, RR_calf, RL_hip, RL_thigh, RL_calf]
|
||||
# Policy 输出顺序取决于 IsaacLab/USD 文件中的定义
|
||||
# 腿部排列共有 4! = 24 种可能性
|
||||
#
|
||||
# 测试组 (腿部顺序映射):
|
||||
# 0: FR-FL-RR-RL (MuJoCo默认顺序)
|
||||
# 1: FR-FL-RL-RR
|
||||
# 2: FR-RR-FL-RL
|
||||
# 3: FR-RR-RL-FL
|
||||
# 4: FR-RL-FL-RR
|
||||
# 5: FR-RL-RR-FL
|
||||
# 6: FL-FR-RR-RL
|
||||
# 7: FL-FR-RL-RR
|
||||
# 8: FL-RR-FR-RL
|
||||
# 9: FL-RR-RL-FR
|
||||
# 10: FL-RL-FR-RR
|
||||
# 11: FL-RL-RR-FR
|
||||
# 12: RR-FR-FL-RL
|
||||
# 13: RR-FR-RL-FL 。。
|
||||
# 14: RR-FL-FR-RL
|
||||
# 15: RR-FL-RL-FR
|
||||
# 16: RR-RL-FR-FL
|
||||
# 17: RR-RL-FL-FR
|
||||
# 18: RL-FR-FL-RR 。。
|
||||
# 19: RL-FR-RR-FL 。。
|
||||
# 20: RL-FL-FR-RR
|
||||
# 21: RL-FL-RR-FR
|
||||
# 22: RL-RR-FR-FL
|
||||
# 23: RL-RR-FL-FR 。。
|
||||
ACTION_MAPPING = 0 # 选择测试组 (0-23)
|
||||
|
||||
# 腿部顺序定义
|
||||
LEG_MAPPING = [
|
||||
[0, 1, 2, 3], # 0: FR-FL-RR-RL (MuJoCo顺序)
|
||||
[0, 1, 3, 2], # 1: FR-FL-RL-RR
|
||||
[0, 2, 1, 3], # 2: FR-RR-FL-RL
|
||||
[0, 2, 3, 1], # 3: FR-RR-RL-FL
|
||||
[0, 3, 1, 2], # 4: FR-RL-FL-RR
|
||||
[0, 3, 2, 1], # 5: FR-RL-RR-FL
|
||||
[1, 0, 2, 3], # 6: FL-FR-RR-RL
|
||||
[1, 0, 3, 2], # 7: FL-FR-RL-RR
|
||||
[1, 2, 0, 3], # 8: FL-RR-FR-RL
|
||||
[1, 2, 3, 0], # 9: FL-RR-RL-FR
|
||||
[1, 3, 0, 2], # 10: FL-RL-FR-RR
|
||||
[1, 3, 2, 0], # 11: FL-RL-RR-FR
|
||||
[2, 0, 1, 3], # 12: RR-FR-FL-RL
|
||||
[2, 0, 3, 1], # 13: RR-FR-RL-FL
|
||||
[2, 1, 0, 3], # 14: RR-FL-FR-RL
|
||||
[2, 1, 3, 0], # 15: RR-FL-RL-FR
|
||||
[2, 3, 0, 1], # 16: RR-RL-FR-FL
|
||||
[2, 3, 1, 0], # 17: RR-RL-FL-FR
|
||||
[3, 0, 1, 2], # 18: RL-FR-FL-RR
|
||||
[3, 0, 2, 1], # 19: RL-FR-RR-FL
|
||||
[3, 1, 0, 2], # 20: RL-FL-FR-RR
|
||||
[3, 1, 2, 0], # 21: RL-FL-RR-FR
|
||||
[3, 2, 0, 1], # 22: RL-RR-FR-FL
|
||||
[3, 2, 1, 0], # 23: RL-RR-FL-FR
|
||||
]
|
||||
|
||||
# 根据选择生成 DOF_MAPPING
|
||||
# 每个腿有3个关节(hip, thigh, calf)
|
||||
leg_order = LEG_MAPPING[ACTION_MAPPING]
|
||||
DOF_MAPPING = np.array([
|
||||
leg_order[0]*3 + 0, leg_order[0]*3 + 1, leg_order[0]*3 + 2, # 腿0
|
||||
leg_order[1]*3 + 0, leg_order[1]*3 + 1, leg_order[1]*3 + 2, # 腿1
|
||||
leg_order[2]*3 + 0, leg_order[2]*3 + 1, leg_order[2]*3 + 2, # 腿2
|
||||
leg_order[3]*3 + 0, leg_order[3]*3 + 1, leg_order[3]*3 + 2, # 腿3
|
||||
])
|
||||
|
||||
# 速度命令范围
|
||||
MAX_LIN_VEL = 1.0 # 最大线速度 m/s
|
||||
MAX_ANG_VEL = 1.0 # 最大角速度 rad/s
|
||||
|
||||
# ============================================================
|
||||
# 缩放参数
|
||||
# ============================================================
|
||||
ACTION_SCALE = 0.15 # RL 输出动作的缩放 (映射到关节位置范围)
|
||||
POS_SCALE = -1.0 # MLP 位置误差缩放
|
||||
VEL_SCALE = 1.0 # MLP 速度缩放
|
||||
TORQUE_SCALE = 1.0 # MLP torque 输出缩放
|
||||
TORQUE_LIMIT = 23.7 # torque 限幅 (N·m)
|
||||
|
||||
# MuJoCo 摩擦系数配置
|
||||
# friction = [滑动摩擦(X), 滑动摩擦(Y), 扭转摩擦]
|
||||
# - friction[0]: X方向滑动摩擦系数
|
||||
# - friction[1]: Y方向滑动摩擦系数 (各向异性材料用)
|
||||
# - friction[2]: 扭转摩擦系数 (绕法线的旋转阻力)
|
||||
#
|
||||
# 常见地面摩擦参数参考:
|
||||
# ┌──────────┬───────────────────┐
|
||||
# │ 地面类型 │ FLOOR_FRICTION │
|
||||
# ├──────────┼───────────────────┤
|
||||
# │ 冰面 │ [0.05, 0.02, 0.01]│
|
||||
# │ 木地板 │ [0.4, 0.3, 0.2 ]│
|
||||
# │ 瓷砖 │ [0.6, 0.4, 0.3 ]│
|
||||
# │ 橡胶垫 │ [1.5, 1.0, 0.5 ]│
|
||||
# │ 粗糙地面 │ [2.0, 1.5, 1.0 ]│
|
||||
# └──────────┴───────────────────┘
|
||||
# 机器人 geom 摩擦系数
|
||||
FRICTION = [0.6, 0.3, 0.3]
|
||||
# 地面摩擦系数 (单独设置)
|
||||
FLOOR_FRICTION = [2.0, 1.5, 1.0]
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 键盘输入读取 (使用后台线程，无需 root 权限)
|
||||
# ============================================================
|
||||
class KeyboardReader:
|
||||
"""使用后台线程读取键盘输入，不需要 root 权限"""
|
||||
|
||||
def __init__(self):
|
||||
self.old_settings = None
|
||||
self.keys_pressed = set()
|
||||
self.running = True
|
||||
self.thread = None
|
||||
|
||||
def init(self):
|
||||
"""初始化终端并启动键盘读取线程"""
|
||||
self.old_settings = termios.tcgetattr(sys.stdin)
|
||||
tty.setcbreak(sys.stdin.fileno())
|
||||
self.thread = threading.Thread(target=self._read_loop, daemon=True)
|
||||
self.thread.start()
|
||||
|
||||
def _read_loop(self):
|
||||
"""后台线程：持续读取键盘输入"""
|
||||
try:
|
||||
while self.running:
|
||||
try:
|
||||
import select
|
||||
if select.select([sys.stdin], [], [], 0.1)[0]:
|
||||
ch = sys.stdin.read(1)
|
||||
if ch:
|
||||
if ch == '\x1b': # ESC
|
||||
self.keys_pressed.add('\x1b')
|
||||
elif ch == ' ':
|
||||
self.keys_pressed.add('space')
|
||||
else:
|
||||
self.keys_pressed.add(ch.lower())
|
||||
except:
|
||||
pass
|
||||
except:
|
||||
pass
|
||||
|
||||
def restore(self):
|
||||
"""恢复终端设置"""
|
||||
self.running = False
|
||||
if self.thread:
|
||||
self.thread.join(timeout=1.0)
|
||||
if self.old_settings:
|
||||
termios.tcsetattr(sys.stdin, termios.TCSADRAIN, self.old_settings)
|
||||
|
||||
def is_key_pressed(self, key):
|
||||
"""检查按键是否被按下"""
|
||||
if key == '\x1b':
|
||||
return '\x1b' in self.keys_pressed
|
||||
elif key == ' ':
|
||||
return 'space' in self.keys_pressed
|
||||
return key in self.keys_pressed
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 辅助函数
|
||||
# ============================================================
|
||||
|
||||
def get_joint_state(model, data, joint_name):
|
||||
"""通过关节名获取关节位置"""
|
||||
joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, joint_name)
|
||||
if joint_id < 0:
|
||||
return None
|
||||
qposadr = model.jnt_qposadr[joint_id]
|
||||
return data.qpos[qposadr]
|
||||
|
||||
|
||||
def compute_observations(data, prev_action, command, default_joint_pos):
|
||||
"""
|
||||
计算 RL policy 观测向量 (48维)
|
||||
|
||||
对应 IsaacLab 观测配置:
|
||||
[0:3] base_lin_vel - 躯干坐标系线速度
|
||||
[3:6] base_ang_vel - 躯干坐标系角速度
|
||||
[6:9] projected_gravity - 重力投影到躯干坐标系
|
||||
[9:12] velocity_commands - 期望速度
|
||||
[12:24] joint_pos_rel - 相对关节位置 (相对于默认位置)
|
||||
[24:36] joint_vel_rel - 相对关节速度 (相对于默认速度)
|
||||
[36:48] prev_actions - 上一步动作
|
||||
"""
|
||||
obs = np.zeros(OBS_DIM, dtype=np.float32)
|
||||
|
||||
# 获取躯干四元数 (MuJoCo 格式: qx, qy, qz, qw)
|
||||
quat = data.qpos[3:7] # (x, y, z, w)
|
||||
|
||||
# 计算旋转矩阵 (从世界坐标系到躯干坐标系)
|
||||
R = quaternion_to_rotation_matrix(quat).T # 转置得到逆变换 (世界->躯干)
|
||||
|
||||
# 1. 躯干线速度 - 转换到躯干坐标系
|
||||
# IsaacLab: root_lin_vel_b = R @ root_lin_vel_w
|
||||
global_lin_vel = data.qvel[0:3]
|
||||
obs[0:3] = R @ global_lin_vel
|
||||
|
||||
# 2. 躯干角速度 - 转换到躯干坐标系
|
||||
# IsaacLab: root_ang_vel_b = R @ root_ang_vel_w
|
||||
global_ang_vel = data.qvel[3:6]
|
||||
obs[3:6] = R @ global_ang_vel
|
||||
|
||||
# 3. 重力投影到躯干坐标系
|
||||
# IsaacLab: projected_gravity_b = R @ gravity_world
|
||||
# 重力向量 (0, 0, -9.81) 在世界坐标系中
|
||||
gravity_world = np.array([0.0, 0.0, -9.81])
|
||||
obs[6:9] = R @ gravity_world
|
||||
|
||||
# 4. 速度命令
|
||||
obs[9:12] = command
|
||||
|
||||
# 5. 相对关节位置 (joint_pos - default_joint_pos)
|
||||
# IsaacLab: joint_pos_rel = joint_pos - default_joint_pos
|
||||
current_joint_pos = data.qpos[7:19]
|
||||
obs[12:24] = current_joint_pos - default_joint_pos
|
||||
|
||||
# 6. 相对关节速度
|
||||
# IsaacLab: joint_vel_rel = joint_vel - default_joint_vel (default = 0)
|
||||
default_joint_vel = np.zeros(12) # IsaacLab 默认关节速度为 0
|
||||
obs[24:36] = data.qvel[6:18] - default_joint_vel
|
||||
|
||||
# 7. 上一步动作
|
||||
obs[36:48] = prev_action
|
||||
|
||||
return obs
|
||||
|
||||
|
||||
def quaternion_to_rotation_matrix(q):
|
||||
"""四元数转旋转矩阵 (MuJoCo 格式: qx, qy, qz, qw)
|
||||
|
||||
返回的矩阵将世界坐标系中的向量转换到以该四元数表示的局部坐标系中
|
||||
"""
|
||||
qx, qy, qz, qw = q
|
||||
# 归一化确保数值稳定性
|
||||
norm = np.sqrt(qx**2 + qy**2 + qz**2 + qw**2)
|
||||
qx, qy, qz, qw = qx/norm, qy/norm, qz/norm, qw/norm
|
||||
return np.array([
|
||||
[1-2*(qy**2+qz**2), 2*(qx*qy-qz*qw), 2*(qx*qz+qy*qw)],
|
||||
[2*(qx*qy+qz*qw), 1-2*(qx**2+qz**2), 2*(qy*qz-qx*qw)],
|
||||
[2*(qx*qz-qy*qw), 2*(qy*qz+qx*qw), 1-2*(qx**2+qy**2)]
|
||||
])
|
||||
|
||||
|
||||
# ============================================================
|
||||
# MLP Actuator Model (从 IsaacLab 迁移)
|
||||
# ============================================================
|
||||
class ActuatorMLP:
|
||||
"""ActuatorNetMLP 模型 - 用于将 position target 转换为 torque
|
||||
|
||||
参考: IsaacLab source/isaaclab/isaaclab/actuators/actuator_net.py
|
||||
使用批量处理，所有 12 个关节一次前向传播
|
||||
"""
|
||||
def __init__(self, model_path, num_joints=12):
|
||||
import torch
|
||||
|
||||
self.num_joints = num_joints
|
||||
|
||||
# 强制使用 CPU 避免 CUDA 版本问题
|
||||
self.network = torch.jit.load(model_path, map_location='cpu')
|
||||
self.network.eval()
|
||||
|
||||
# GO1 MLP 配置 (来自 GO1_ACTUATOR_CFG)
|
||||
self.pos_scale = POS_SCALE
|
||||
self.vel_scale = VEL_SCALE
|
||||
self.torque_scale = TORQUE_SCALE
|
||||
self.input_idx = [0, 1, 2] # 3步历史
|
||||
self.input_order = "pos_vel"
|
||||
|
||||
# 历史缓冲区 (每个关节独立)
|
||||
history_len = max(self.input_idx) + 1 # = 3
|
||||
self.pos_error_history = np.zeros((history_len, num_joints), dtype=np.float32)
|
||||
self.vel_history = np.zeros((history_len, num_joints), dtype=np.float32)
|
||||
|
||||
# Torque 限幅 (来自 URDF 和 GO1_ACTUATOR_CFG)
|
||||
self.effort_limit = TORQUE_LIMIT # N·m
|
||||
|
||||
def reset(self):
|
||||
"""重置历史缓冲区"""
|
||||
self.pos_error_history.fill(0.0)
|
||||
self.vel_history.fill(0.0)
|
||||
|
||||
def compute_torque(self, pos_target, current_pos, current_vel):
|
||||
"""计算 torque (使用批量处理，匹配 IsaacLab 实现)
|
||||
|
||||
Args:
|
||||
pos_target: 目标位置 (12,)
|
||||
current_pos: 当前关节位置 (12,)
|
||||
current_vel: 当前关节速度 (12,)
|
||||
|
||||
Returns:
|
||||
torque: 力矩命令 (12,)
|
||||
"""
|
||||
import torch
|
||||
|
||||
# 1. 计算 position error
|
||||
pos_error = pos_target - current_pos # (12,)
|
||||
|
||||
# 2. 更新历史 (移动队列)
|
||||
self.pos_error_history = np.roll(self.pos_error_history, 1, axis=0)
|
||||
self.pos_error_history[0] = pos_error
|
||||
|
||||
self.vel_history = np.roll(self.vel_history, 1, axis=0)
|
||||
self.vel_history[0] = current_vel
|
||||
|
||||
# 3. 批量计算 torque (IsaacLab 风格，一次性处理所有关节)
|
||||
# pos_error_history: (history_len, num_joints) = (3, 12)
|
||||
# vel_history: (history_len, num_joints) = (3, 12)
|
||||
#
|
||||
# IsaacLab 构建输入的方式:
|
||||
# pos_input = torch.cat([history[:, i].unsqueeze(2) for i in input_idx], dim=2) # -> (num_envs, num_joints, 3)
|
||||
# pos_input = pos_input.view(num_envs * num_joints, -1) # -> (12, 3)
|
||||
# 然后 concat pos_input 和 vel_input 得到 (12, 6)
|
||||
|
||||
with torch.inference_mode():
|
||||
# 构建 pos_input: (num_joints, 3) = (12, 3)
|
||||
# 从 history 矩阵中提取 input_idx 指定的行，然后转置
|
||||
# pos_error_history[self.input_idx, :] shape = (3, 12)
|
||||
# 转置后变成 (12, 3)，每行是一个关节的 3 步位置误差历史
|
||||
pos_input = self.pos_error_history[self.input_idx, :].T * self.pos_scale # (12, 3)
|
||||
|
||||
# 构建 vel_input: (num_joints, 3) = (12, 3)
|
||||
vel_input = self.vel_history[self.input_idx, :].T * self.vel_scale # (12, 3)
|
||||
|
||||
# 合并为网络输入: (12, 6)
|
||||
if self.input_order == "pos_vel":
|
||||
network_input = np.concatenate([pos_input, vel_input], axis=1)
|
||||
else:
|
||||
network_input = np.concatenate([vel_input, pos_input], axis=1)
|
||||
|
||||
# 批量推理 (12, 6) -> (12, 1)
|
||||
input_tensor = torch.from_numpy(network_input).float() # (12, 6)
|
||||
torque_output = self.network(input_tensor).numpy().flatten() # (12,)
|
||||
|
||||
# 限幅
|
||||
torques = np.clip(torque_output * self.torque_scale,
|
||||
-self.effort_limit, self.effort_limit)
|
||||
|
||||
return torques.astype(np.float32)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 主程序
|
||||
# ============================================================
|
||||
|
||||
def main():
|
||||
# 1. 加载 MuJoCo 模型
|
||||
os.chdir(MESH_PATH)
|
||||
with open(XML_PATH, 'r') as f:
|
||||
xml_content = f.read()
|
||||
xml_content = xml_content.replace('meshdir="../meshes/"', f'meshdir="{MESH_PATH}"')
|
||||
xml_content = xml_content.replace('<default>\n=', '<default>\n')
|
||||
|
||||
model = mujoco.MjModel.from_xml_string(xml_content)
|
||||
data = mujoco.MjData(model)
|
||||
|
||||
print(f"[INFO] MuJoCo Model: {model.nbody} bodies, {model.nq} DoF, {model.nu} actuators")
|
||||
print(f"[INFO] 关节顺序: {[mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i) for i in range(1, 13)]}")
|
||||
|
||||
# 2. 设置摩擦系数
|
||||
floor_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "floor")
|
||||
if floor_id >= 0:
|
||||
model.geom_friction[floor_id] = FLOOR_FRICTION
|
||||
print(f"[INFO] 地面摩擦系数设置为: {FLOOR_FRICTION}")
|
||||
else:
|
||||
print(f"[WARNING] 未找到 floor geom，使用默认摩擦系数")
|
||||
for i in range(model.ngeom):
|
||||
if i != floor_id:
|
||||
model.geom_friction[i] = FRICTION
|
||||
print(f"[INFO] 机器人摩擦系数设置为: {FRICTION}")
|
||||
|
||||
# 3. 加载 ONNX 策略
|
||||
import onnxruntime as ort
|
||||
sess = ort.InferenceSession(ONNX_PATH, providers=['CPUExecutionProvider'])
|
||||
print(f"[INFO] ONNX Policy loaded from: {ONNX_PATH}")
|
||||
print(f"[INFO] 观测维度: {OBS_DIM}, 动作维度: {ACTION_DIM}")
|
||||
|
||||
# 3. 加载 MLP Actuator 模型 (IsaacLab 风格)
|
||||
try:
|
||||
actuator_mlp = ActuatorMLP(MLP_PATH, num_joints=ACTION_DIM)
|
||||
print(f"[INFO] ActuatorMLP loaded from: {MLP_PATH}")
|
||||
use_mlp = True
|
||||
except Exception as e:
|
||||
print(f"[WARNING] Failed to load ActuatorMLP: {e}")
|
||||
print(f"[INFO] Falling back to PD control")
|
||||
use_mlp = False
|
||||
|
||||
# 4. 初始化
|
||||
obs = np.zeros(OBS_DIM, dtype=np.float32)
|
||||
prev_action = np.zeros(ACTION_DIM, dtype=np.float32)
|
||||
|
||||
# IsaacLab UNITREE_GO1_CFG 默认关节位置 (用于计算相对关节位置)
|
||||
# 关节顺序: FR_hip, FR_thigh, FR_calf, FL_hip, FL_thigh, FL_calf, RR_hip, RR_thigh, RR_calf, RL_hip, RL_thigh, RL_calf
|
||||
default_joint_pos = np.array([
|
||||
-0.1, # FR_hip
|
||||
0.8, # FR_thigh
|
||||
-1.5, # FR_calf
|
||||
0.1, # FL_hip
|
||||
0.8, # FL_thigh
|
||||
-1.5, # FL_calf
|
||||
-0.1, # RR_hip
|
||||
1.0, # RR_thigh
|
||||
-1.5, # RR_calf
|
||||
0.1, # RL_hip
|
||||
1.0, # RL_thigh
|
||||
-1.5, # RL_calf
|
||||
])
|
||||
|
||||
# 重置机器人关节位置 (与 default_joint_pos 一致)
|
||||
# 注意: MuJoCo 关节角度单位是 rad
|
||||
# URDF 关节范围:
|
||||
# Hip: -0.803 ~ 0.803 rad
|
||||
# Thigh: -1.047 ~ 4.189 rad
|
||||
# Calf: -2.697 ~ -0.916 rad
|
||||
crouch_pos = np.array([
|
||||
-0.1, # FR_hip (URDF: -0.803 ~ 0.803)
|
||||
0.8, # FR_thigh (URDF: -1.047 ~ 4.189)
|
||||
-1.5, # FR_calf (URDF: -2.697 ~ -0.916)
|
||||
0.1, # FL_hip
|
||||
0.8, # FL_thigh
|
||||
-1.5, # FL_calf
|
||||
-0.1, # RR_hip
|
||||
1.0, # RR_thigh
|
||||
-1.5, # RR_calf
|
||||
0.1, # RL_hip
|
||||
1.0, # RL_thigh
|
||||
-1.5, # RL_calf
|
||||
])
|
||||
data.qpos[7:19] = crouch_pos
|
||||
# 重置 qvel 为零，避免初始速度导致不稳定
|
||||
data.qvel[:] = 0.0
|
||||
# 将躯干抬高一点，避免初始碰撞
|
||||
data.qpos[2] = 0.35
|
||||
mujoco.mj_step(model, data) # 执行一步让状态更新
|
||||
|
||||
# 重置 MLP actuator 历史
|
||||
if use_mlp:
|
||||
actuator_mlp.reset()
|
||||
|
||||
print("[INFO] 机器人初始化完成，开始 RL 策略推理...", flush=True)
|
||||
|
||||
# 初始化键盘控制
|
||||
keyboard_reader = KeyboardReader()
|
||||
keyboard_reader.init()
|
||||
print("[INFO] 键盘控制已启用!", flush=True)
|
||||
print(" W/S: 前进/后退", flush=True)
|
||||
print(" A/D: 左转/右转", flush=True)
|
||||
print(" Q/E: 侧向左移/右移", flush=True)
|
||||
print(" 空格: 停止", flush=True)
|
||||
print(" ESC: 退出", flush=True)
|
||||
|
||||
# 4. 交互式可视化
|
||||
import mujoco.viewer as mv
|
||||
view = mv.launch_passive(model, data)
|
||||
print("[INFO] 交互式查看器已启动!", flush=True)
|
||||
print(" 鼠标拖拽旋转视角, 滚轮缩放, ESC退出", flush=True)
|
||||
|
||||
step_count = 0
|
||||
last_inference_time = time.time()
|
||||
inference_interval = 0.02 # 50Hz = 20ms
|
||||
torque = np.zeros(ACTION_DIM, dtype=np.float32) # 初始化 torque
|
||||
import sys
|
||||
sys.stderr.write(f"[DEBUG] Starting loop\n")
|
||||
sys.stderr.flush()
|
||||
|
||||
# 命令速度 (使用 numpy array 以便在循环中修改)
|
||||
command = np.array([1.0, 0.0, 0.0], dtype=np.float32) # vx, vy, yaw_rate
|
||||
|
||||
while view.is_running() and not g_exit_requested:
|
||||
current_time = time.time()
|
||||
|
||||
# W/S: 前进/后退 (vx)
|
||||
if keyboard_reader.is_key_pressed('w'):
|
||||
command[0] = MAX_LIN_VEL
|
||||
elif keyboard_reader.is_key_pressed('s'):
|
||||
command[0] = -MAX_LIN_VEL
|
||||
else:
|
||||
command[0] = 0.0
|
||||
|
||||
# Q/E: 侧向移动 (vy)
|
||||
if keyboard_reader.is_key_pressed('q'):
|
||||
command[1] = MAX_LIN_VEL
|
||||
elif keyboard_reader.is_key_pressed('e'):
|
||||
command[1] = -MAX_LIN_VEL
|
||||
else:
|
||||
command[1] = 0.0
|
||||
|
||||
# A/D: 左转/右转 (yaw_rate)
|
||||
if keyboard_reader.is_key_pressed('a'):
|
||||
command[2] = MAX_ANG_VEL
|
||||
elif keyboard_reader.is_key_pressed('d'):
|
||||
command[2] = -MAX_ANG_VEL
|
||||
else:
|
||||
command[2] = 0.0
|
||||
|
||||
# 空格: 停止
|
||||
if keyboard_reader.is_key_pressed(' '):
|
||||
command[:] = 0.0
|
||||
|
||||
# ESC: 退出
|
||||
if keyboard_reader.is_key_pressed('\x1b'):
|
||||
sys.stderr.write("[INFO] ESC pressed, exiting...\n")
|
||||
break
|
||||
|
||||
# 50Hz 频率限制推理
|
||||
if current_time - last_inference_time >= inference_interval:
|
||||
# 计算观测
|
||||
obs = compute_observations(data, prev_action, command, default_joint_pos)
|
||||
|
||||
# 推理获取动作
|
||||
action = sess.run(None, {'obs': obs.reshape(1, -1)})[0][0]
|
||||
|
||||
# 动作缩放 (RL 输出是位置目标, 范围 [-1, 1], 映射到关节范围)
|
||||
# IsaacLab 配置: actions.joint_pos.scale = 0.25
|
||||
action_scaled = action * ACTION_SCALE
|
||||
|
||||
# 应用关节顺序映射 (policy输出 -> MuJoCo关节)
|
||||
action_scaled_mapped = action_scaled[DOF_MAPPING]
|
||||
|
||||
# 获取当前关节状态
|
||||
current_pos = data.qpos[7:19]
|
||||
current_vel = data.qvel[6:18]
|
||||
|
||||
# 使用 MLP 计算 torque (与 IsaacLab 一致)
|
||||
if use_mlp:
|
||||
torque = actuator_mlp.compute_torque(
|
||||
action_scaled_mapped, current_pos, current_vel
|
||||
)
|
||||
else:
|
||||
# Fallback PD 控制
|
||||
kp = 10.0
|
||||
kd = 0.0
|
||||
torque = kp * (action_scaled_mapped - current_pos) - kd * current_vel
|
||||
|
||||
# 保存上一步动作
|
||||
prev_action = action.copy()
|
||||
last_inference_time = current_time
|
||||
|
||||
# 应用控制 (力矩控制) - 每步都应用
|
||||
data.ctrl[:] = torque
|
||||
|
||||
# 执行仿真
|
||||
mujoco.mj_step(model, data)
|
||||
view.sync()
|
||||
|
||||
step_count += 1
|
||||
trunk_z = data.qpos[2]
|
||||
lin_vel = np.linalg.norm(data.qvel[0:3])
|
||||
ang_vel = np.linalg.norm(data.qvel[3:6])
|
||||
sys.stderr.write(f" Step {step_count}: cmd=[{command[0]:.2f}, {command[1]:.2f}, {command[2]:.2f}] | "
|
||||
f"trunk_z={trunk_z:.3f}m, vel={lin_vel:.3f}m/s\n")
|
||||
sys.stderr.flush()
|
||||
|
||||
# 恢复终端设置
|
||||
if g_exit_requested:
|
||||
sys.stderr.write("[INFO] Ctrl+C pressed, exiting...\n")
|
||||
keyboard_reader.restore()
|
||||
view.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,470 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Go1 RL Policy 推理 - 移植自 unitree_rl 项目
|
||||
基于 https://github.com/dstx123/unitree_rl
|
||||
|
||||
模型: body.jit + adapt.jit (GRU-based policy with history)
|
||||
训练环境: IsaacGym
|
||||
|
||||
使用方法:
|
||||
python3 /home/8x54zj-m/unitree_mujoco/go1_unitree_rl_inference.py
|
||||
|
||||
键盘控制:
|
||||
W/S: 前进/后退
|
||||
A/D: 左转/右转
|
||||
Q/E: 侧向左移/右移
|
||||
空格: 停止
|
||||
ESC: 退出
|
||||
"""
|
||||
import numpy as np
|
||||
import mujoco
|
||||
from mujoco import viewer
|
||||
import os
|
||||
import sys
|
||||
import tty
|
||||
import termios
|
||||
import signal
|
||||
import threading
|
||||
import time
|
||||
import torch
|
||||
|
||||
# ============================================================
|
||||
# 全局退出标志
|
||||
# ============================================================
|
||||
g_exit_requested = False
|
||||
|
||||
def signal_handler(signum, frame):
|
||||
global g_exit_requested
|
||||
g_exit_requested = True
|
||||
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
# ============================================================
|
||||
# 配置
|
||||
# ============================================================
|
||||
BODY_MODEL_PATH = "/tmp/unitree_rl/src/unitree_guide/unitree_guide/model/body.jit"
|
||||
ADAPT_MODEL_PATH = "/tmp/unitree_rl/src/unitree_guide/unitree_guide/model/adapt.jit"
|
||||
XML_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/xml/go1.xml"
|
||||
MESH_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/meshes"
|
||||
|
||||
# ============================================================
|
||||
# 模型参数 (来自 unitree_rl/State_RL.h)
|
||||
# ============================================================
|
||||
NUM_OBS = 45 # 观测维度
|
||||
NUM_OBS_HISTORY = 10 # 历史观测步数
|
||||
OBS_BUFFER_SIZE = NUM_OBS * NUM_OBS_HISTORY # 450
|
||||
|
||||
# 观测缩放
|
||||
SCALE_LIN_VEL = 2.0
|
||||
SCALE_ANG_VEL = 0.25
|
||||
SCALE_COMMANDS = np.array([SCALE_LIN_VEL, SCALE_LIN_VEL, SCALE_ANG_VEL])
|
||||
SCALE_DOF_POS = 1.0
|
||||
SCALE_DOF_VEL = 0.05
|
||||
|
||||
# 动作缩放
|
||||
ACTION_SCALE = 0.25
|
||||
HIP_SCALE_REDUCTION = 0.5 # hip关节额外缩放
|
||||
|
||||
# 限幅
|
||||
CLIP_OBSERVATIONS = 100.0
|
||||
CLIP_ACTIONS = 100.0
|
||||
|
||||
# 默认关节角度 (与 unitree_rl 一致)
|
||||
DEFAULT_DOF_POS = np.array([
|
||||
-0.1, 0.8, -1.5, # FR: hip, thigh, calf
|
||||
0.1, 0.8, -1.5, # FL: hip, thigh, calf
|
||||
-0.1, 1.0, -1.5, # RR: hip, thigh, calf
|
||||
0.1, 1.0, -1.5 # RL: hip, thigh, calf
|
||||
])
|
||||
|
||||
# Motor order in code: FR(3,4,5), FL(0,1,2), RR(9,10,11), RL(6,7,8)
|
||||
# Model order: FL(0,1,2), FR(3,4,5), RL(6,7,8), RR(9,10,11)
|
||||
# dof_mapping = {3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8}
|
||||
# Inverse mapping (model idx -> motor idx): 3->0, 4->1, 5->2, 0->3, 1->4, 2->5, 9->6, 10->7, 11->8, 6->9, 7->10, 8->11
|
||||
# This maps: model_output[0..11] -> motor_q[0..11]
|
||||
DOF_MAPPING = np.array([3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8])
|
||||
|
||||
# 速度命令范围
|
||||
MAX_LIN_VEL = 1.0
|
||||
MAX_ANG_VEL = 1.0
|
||||
|
||||
# PD控制参数 (来自 unitree_rl)
|
||||
KP = 7.0
|
||||
KD = 1.0
|
||||
|
||||
# MuJoCo 摩擦系数配置
|
||||
# friction = [滑动摩擦(X), 滑动摩擦(Y), 扭转摩擦]
|
||||
# - friction[0]: X方向滑动摩擦系数
|
||||
# - friction[1]: Y方向滑动摩擦系数 (各向异性材料用)
|
||||
# - friction[2]: 扭转摩擦系数 (绕法线的旋转阻力)
|
||||
#
|
||||
# 常见地面摩擦参数参考:
|
||||
# ┌──────────┬───────────────────┐
|
||||
# │ 地面类型 │ FLOOR_FRICTION │
|
||||
# ├──────────┼───────────────────┤
|
||||
# │ 冰面 │ [0.05, 0.02, 0.01]│
|
||||
# │ 木地板 │ [0.4, 0.3, 0.2 ]│
|
||||
# │ 瓷砖 │ [0.6, 0.4, 0.3 ]│
|
||||
# │ 橡胶垫 │ [1.5, 1.0, 0.5 ]│
|
||||
# │ 粗糙地面 │ [2.0, 1.5, 1.0 ]│
|
||||
# └──────────┴───────────────────┘
|
||||
# 机器人 geom 摩擦系数
|
||||
FRICTION = [0.6, 0.3, 0.3]
|
||||
# 地面摩擦系数 (单独设置)
|
||||
FLOOR_FRICTION = [2.0, 1.5, 1.0]
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 键盘输入读取 (termios 非阻塞方式，无 root 依赖)
|
||||
# ============================================================
|
||||
import select
|
||||
|
||||
class KeyboardReader:
|
||||
def __init__(self):
|
||||
self.keys_pressed = set()
|
||||
self.last_key_time = 0
|
||||
self.timeout = 0.1 # 无按键超时时间（秒），超时后清除所有按键
|
||||
|
||||
def init(self):
|
||||
import termios
|
||||
self.old_settings = termios.tcgetattr(sys.stdin)
|
||||
# 非canonical模式，关闭回显
|
||||
new_settings = termios.tcgetattr(sys.stdin)
|
||||
new_settings[3] = new_settings[3] & ~termios.ICANON & ~termios.ECHO
|
||||
new_settings[6][termios.VMIN] = 0 # 非阻塞
|
||||
new_settings[6][termios.VTIME] = 0
|
||||
termios.tcsetattr(sys.stdin, termios.TCSADRAIN, new_settings)
|
||||
print("[INFO] Keyboard reader initialized (termios mode, no root required)")
|
||||
|
||||
def _read_key(self):
|
||||
"""尝试读取一个按键，返回按键字符或None"""
|
||||
try:
|
||||
if select.select([sys.stdin], [], [], 0)[0]:
|
||||
ch = sys.stdin.read(1)
|
||||
return ch
|
||||
except:
|
||||
pass
|
||||
return None
|
||||
|
||||
def update(self):
|
||||
"""在主循环中调用，更新按键状态"""
|
||||
import time
|
||||
current_time = time.time()
|
||||
|
||||
# 尝试读取按键
|
||||
ch = self._read_key()
|
||||
if ch:
|
||||
if ch == '\x1b': # ESC
|
||||
self.keys_pressed.add('escape')
|
||||
elif ch == ' ':
|
||||
self.keys_pressed.add('space')
|
||||
elif ch == '\n' or ch == '\r':
|
||||
pass
|
||||
else:
|
||||
self.keys_pressed.add(ch.lower())
|
||||
self.last_key_time = current_time
|
||||
elif current_time - self.last_key_time > self.timeout:
|
||||
# 超时后清除所有按键（模拟按键释放）
|
||||
self.keys_pressed.clear()
|
||||
|
||||
def restore(self):
|
||||
import termios
|
||||
if hasattr(self, 'old_settings'):
|
||||
termios.tcsetattr(sys.stdin, termios.TCSADRAIN, self.old_settings)
|
||||
|
||||
def is_key_pressed(self, key):
|
||||
if key == '\x1b' or key == 'escape':
|
||||
return 'escape' in self.keys_pressed
|
||||
elif key == ' ':
|
||||
return 'space' in self.keys_pressed
|
||||
return key.lower() in self.keys_pressed
|
||||
|
||||
def clear(self):
|
||||
self.keys_pressed.clear()
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 辅助函数
|
||||
# ============================================================
|
||||
|
||||
def quaternion_to_rotation_matrix(q):
|
||||
"""四元数转旋转矩阵 (MuJoCo 格式: qx, qy, qz, qw)"""
|
||||
qx, qy, qz, qw = q
|
||||
norm = np.sqrt(qx**2 + qy**2 + qz**2 + qw**2)
|
||||
qx, qy, qz, qw = qx/norm, qy/norm, qz/norm, qw/norm
|
||||
return np.array([
|
||||
[1-2*(qy**2+qz**2), 2*(qx*qy-qz*qw), 2*(qx*qz+qy*qw)],
|
||||
[2*(qx*qy+qz*qw), 1-2*(qx**2+qz**2), 2*(qy*qz-qx*qw)],
|
||||
[2*(qx*qz-qy*qw), 2*(qy*qz+qx*qw), 1-2*(qx**2+qy**2)]
|
||||
])
|
||||
|
||||
|
||||
def quat_rotate_inverse(q, v):
|
||||
"""四元数逆旋转 (世界坐标系 -> 躯干坐标系)
|
||||
|
||||
对应 IsaacLab 的 quat_apply_inverse
|
||||
q: (w, x, y, z) 格式四元数
|
||||
v: (x, y, z) 向量
|
||||
"""
|
||||
q_w, q_x, q_y, q_z = q[3], q[0], q[1], q[2]
|
||||
return np.array([
|
||||
v[0] * (2*q_w**2 - 1) + 2*q_x*(q_y*v[2] - q_z*v[1]) + 2*q_w*(q_z*v[0] - q_x*v[2]),
|
||||
v[1] * (2*q_w**2 - 1) + 2*q_y*(q_z*v[0] - q_x*v[2]) + 2*q_w*(q_x*v[1] - q_y*v[0]),
|
||||
v[2] * (2*q_w**2 - 1) + 2*q_z*(q_x*v[1] - q_y*v[0]) + 2*q_w*(q_y*v[0] - q_x*v[1])
|
||||
])
|
||||
|
||||
|
||||
def compute_observations(data, prev_action, commands, default_dof_pos):
|
||||
"""
|
||||
计算 RL policy 观测向量 (45维)
|
||||
|
||||
对应 unitree_rl 的观测计算:
|
||||
[0:3] body_ang_vel - 躯干角速度 (body frame)
|
||||
[3:6] projected_gravity - 重力投影到躯干坐标系
|
||||
[6:9] commands - 速度命令 (已缩放)
|
||||
[9:21] dof_pos_rel - 相对关节位置
|
||||
[21:33] dof_vel - 关节速度
|
||||
[33:45] prev_action - 上一步动作
|
||||
"""
|
||||
obs = np.zeros(NUM_OBS, dtype=np.float32)
|
||||
|
||||
# 获取四元数 (MuJoCo 格式: x, y, z, w)
|
||||
quat_mujoco = data.qpos[3:7] # (x, y, z, w)
|
||||
# 转换为 unitree_rl 格式 (x, y, z, w) for quat_rotate_inverse
|
||||
|
||||
# 1. 躯干角速度 (body frame)
|
||||
# unitree_rl 直接使用 imu.gyroscope (已是在 body frame)
|
||||
# MuJoCo 的 qvel[3:6] 是世界坐标系，需要转换
|
||||
ang_vel_world = data.qvel[3:6]
|
||||
ang_vel_body = quat_rotate_inverse(np.concatenate([quat_mujoco[3:4], quat_mujoco[0:3]]), ang_vel_world)
|
||||
obs[0:3] = ang_vel_body * SCALE_ANG_VEL
|
||||
|
||||
# 2. 重力投影 (body frame)
|
||||
# gravity_vec = (0, 0, -1) in unitree_rl
|
||||
gravity_world = np.array([0.0, 0.0, -1.0])
|
||||
projected_gravity = quat_rotate_inverse(np.concatenate([quat_mujoco[3:4], quat_mujoco[0:3]]), gravity_world)
|
||||
obs[3:6] = projected_gravity
|
||||
|
||||
# 3. 速度命令 (已缩放)
|
||||
obs[6:9] = commands * SCALE_COMMANDS
|
||||
|
||||
# 4. 相对关节位置
|
||||
# Motor order: FR(3,4,5), FL(0,1,2), RR(9,10,11), RL(6,7,8)
|
||||
# Model order: FL(0,1,2), FR(3,4,5), RL(6,7,8), RR(9,10,11)
|
||||
motor_pos = data.qpos[7:19] # MuJoCo motor order
|
||||
# Map to model order
|
||||
model_order_pos = motor_pos[DOF_MAPPING] # (12,)
|
||||
dof_pos_rel = model_order_pos - default_dof_pos
|
||||
obs[9:21] = dof_pos_rel * SCALE_DOF_POS
|
||||
|
||||
# 5. 关节速度
|
||||
motor_vel = data.qvel[6:18]
|
||||
model_order_vel = motor_vel[DOF_MAPPING]
|
||||
obs[21:33] = model_order_vel * SCALE_DOF_VEL
|
||||
|
||||
# 6. 上一步动作
|
||||
obs[33:45] = prev_action
|
||||
|
||||
# 限幅
|
||||
obs = np.clip(obs, -CLIP_OBSERVATIONS, CLIP_OBSERVATIONS)
|
||||
|
||||
return obs
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 主程序
|
||||
# ============================================================
|
||||
|
||||
def main():
|
||||
# 1. 加载 MuJoCo 模型
|
||||
os.chdir(MESH_PATH)
|
||||
with open(XML_PATH, 'r') as f:
|
||||
xml_content = f.read()
|
||||
xml_content = xml_content.replace('meshdir="../meshes/"', f'meshdir="{MESH_PATH}"')
|
||||
xml_content = xml_content.replace('<default>\n=', '<default>\n')
|
||||
|
||||
model = mujoco.MjModel.from_xml_string(xml_content)
|
||||
data = mujoco.MjData(model)
|
||||
|
||||
print(f"[INFO] MuJoCo Model: {model.nbody} bodies, {model.nq} DoF, {model.nu} actuators")
|
||||
print(f"[INFO] 关节顺序: {[mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i) for i in range(1, 13)]}")
|
||||
|
||||
# 2. 设置摩擦系数
|
||||
# friction = [滑动摩擦, 扭转摩擦, 滚动摩擦]
|
||||
# 设置地面摩擦系数 (通过 geom 名称查找)
|
||||
floor_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "floor")
|
||||
if floor_id >= 0:
|
||||
model.geom_friction[floor_id] = FLOOR_FRICTION
|
||||
print(f"[INFO] 地面摩擦系数设置为: {FLOOR_FRICTION}")
|
||||
else:
|
||||
print(f"[WARNING] 未找到 floor geom，使用默认摩擦系数")
|
||||
|
||||
# 设置机器人 geom 的摩擦系数 (排除地面)
|
||||
for i in range(model.ngeom):
|
||||
if i != floor_id:
|
||||
model.geom_friction[i] = FRICTION
|
||||
print(f"[INFO] 机器人摩擦系数设置为: {FRICTION}")
|
||||
|
||||
# 3. 加载 RL 模型
|
||||
body_module = torch.jit.load(BODY_MODEL_PATH, map_location='cpu')
|
||||
adapt_module = torch.jit.load(ADAPT_MODEL_PATH, map_location='cpu')
|
||||
body_module.eval()
|
||||
adapt_module.eval()
|
||||
print(f"[INFO] Body model loaded from: {BODY_MODEL_PATH}")
|
||||
print(f"[INFO] Adapt model loaded from: {ADAPT_MODEL_PATH}")
|
||||
print(f"[INFO] 观测维度: {NUM_OBS}, 历史步数: {NUM_OBS_HISTORY}")
|
||||
|
||||
# 3. 初始化
|
||||
obs = np.zeros(NUM_OBS, dtype=np.float32)
|
||||
prev_action = np.zeros(12, dtype=np.float32)
|
||||
obs_buffer = np.zeros(OBS_BUFFER_SIZE, dtype=np.float32) # 10 * 45 = 450
|
||||
|
||||
# 初始化机器人姿态
|
||||
data.qpos[7:19] = DEFAULT_DOF_POS
|
||||
data.qvel[:] = 0.0
|
||||
data.qpos[2] = 0.35 # 抬高躯干
|
||||
mujoco.mj_step(model, data)
|
||||
|
||||
print("[INFO] 机器人初始化完成，开始 RL 策略推理...")
|
||||
|
||||
# 4. 初始化键盘控制
|
||||
keyboard_reader = KeyboardReader()
|
||||
keyboard_reader.init()
|
||||
print("[INFO] 键盘控制已启用!")
|
||||
print(" W/S: 前进/后退")
|
||||
print(" A/D: 左转/右转")
|
||||
print(" Q/E: 侧向左移/右移")
|
||||
print(" 空格: 停止")
|
||||
print(" ESC: 退出")
|
||||
|
||||
# 5. 启动交互式查看器
|
||||
import mujoco.viewer as mv
|
||||
view = mv.launch_passive(model, data)
|
||||
print("[INFO] 交互式查看器已启动!")
|
||||
|
||||
# 初始化观测 buffer (填充历史)
|
||||
for _ in range(NUM_OBS_HISTORY):
|
||||
obs = compute_observations(data, prev_action, np.array([0.0, 0.0, 0.0]), DEFAULT_DOF_POS)
|
||||
obs_buffer = np.concatenate([obs_buffer[NUM_OBS:], obs])
|
||||
|
||||
step_count = 0
|
||||
last_inference_time = time.time()
|
||||
inference_interval = 0.02 # 50Hz
|
||||
|
||||
command = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
||||
joint_targets = DEFAULT_DOF_POS.copy() # 初始化关节目标
|
||||
|
||||
while view.is_running() and not g_exit_requested:
|
||||
current_time = time.time()
|
||||
|
||||
# 更新键盘状态
|
||||
keyboard_reader.update()
|
||||
|
||||
# 读取键盘输入 (切换模式: 按下切换状态)
|
||||
# 空格: 停止所有
|
||||
if keyboard_reader.is_key_pressed(' '):
|
||||
command[:] = 0.0
|
||||
keyboard_reader.clear()
|
||||
|
||||
# W/S: 前进/后退 (vx)
|
||||
if keyboard_reader.is_key_pressed('w'):
|
||||
command[0] = MAX_LIN_VEL
|
||||
command[1] = 0.0 # 清除侧向
|
||||
command[2] = 0.0 # 清除转向
|
||||
keyboard_reader.clear()
|
||||
|
||||
if keyboard_reader.is_key_pressed('s'):
|
||||
command[0] = -MAX_LIN_VEL
|
||||
command[1] = 0.0
|
||||
command[2] = 0.0
|
||||
keyboard_reader.clear()
|
||||
|
||||
# Q/E: 侧向移动 (vy)
|
||||
if keyboard_reader.is_key_pressed('q'):
|
||||
command[1] = MAX_LIN_VEL
|
||||
command[0] = 0.0
|
||||
command[2] = 0.0
|
||||
keyboard_reader.clear()
|
||||
|
||||
if keyboard_reader.is_key_pressed('e'):
|
||||
command[1] = -MAX_LIN_VEL
|
||||
command[0] = 0.0
|
||||
command[2] = 0.0
|
||||
keyboard_reader.clear()
|
||||
|
||||
# A/D: 左转/右转 (yaw_rate)
|
||||
if keyboard_reader.is_key_pressed('a'):
|
||||
command[2] = MAX_ANG_VEL
|
||||
command[0] = 0.0
|
||||
command[1] = 0.0
|
||||
keyboard_reader.clear()
|
||||
|
||||
if keyboard_reader.is_key_pressed('d'):
|
||||
command[2] = -MAX_LIN_VEL
|
||||
command[0] = 0.0
|
||||
command[1] = 0.0
|
||||
keyboard_reader.clear()
|
||||
|
||||
if keyboard_reader.is_key_pressed('\x1b'):
|
||||
print("[INFO] ESC pressed, exiting...")
|
||||
break
|
||||
|
||||
# 50Hz 推理
|
||||
if current_time - last_inference_time >= inference_interval:
|
||||
# 计算观测
|
||||
obs = compute_observations(data, prev_action, command, DEFAULT_DOF_POS)
|
||||
|
||||
# 更新观测历史 buffer
|
||||
obs_buffer = np.concatenate([obs_buffer[NUM_OBS:], obs])
|
||||
|
||||
# 模型推理
|
||||
with torch.inference_mode():
|
||||
obs_buffer_tensor = torch.from_numpy(obs_buffer).float().unsqueeze(0) # (1, 450)
|
||||
obs_tensor = torch.from_numpy(obs).float().unsqueeze(0) # (1, 45)
|
||||
|
||||
# Adapt 模块: 编码历史观测为 latent
|
||||
latent = adapt_module(obs_buffer_tensor) # (1, 21)
|
||||
|
||||
# Body 模块: 结合当前观测和 latent 生成动作
|
||||
combined_input = torch.cat([obs_tensor, latent], dim=1) # (1, 66)
|
||||
action = body_module(combined_input).numpy().flatten() # (12,)
|
||||
|
||||
# 动作后处理
|
||||
action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
|
||||
actions_scaled = action * ACTION_SCALE
|
||||
|
||||
# Hip关节额外缩放
|
||||
hip_indices = [0, 3, 6, 9]
|
||||
for i in hip_indices:
|
||||
actions_scaled[i] *= HIP_SCALE_REDUCTION
|
||||
|
||||
# 转换为电机角度目标
|
||||
joint_targets = actions_scaled[DOF_MAPPING] + DEFAULT_DOF_POS
|
||||
|
||||
# 保存上一步动作
|
||||
prev_action = action.copy()
|
||||
last_inference_time = current_time
|
||||
|
||||
# PD 控制 (每步执行)
|
||||
current_pos = data.qpos[7:19]
|
||||
current_vel = data.qvel[6:18]
|
||||
torque = KP * (joint_targets - current_pos) - KD * current_vel
|
||||
|
||||
# 应用控制
|
||||
data.ctrl[:] = torque
|
||||
mujoco.mj_step(model, data)
|
||||
view.sync()
|
||||
|
||||
step_count += 1
|
||||
if step_count % 100 == 0:
|
||||
trunk_z = data.qpos[2]
|
||||
lin_vel = np.linalg.norm(data.qvel[0:3])
|
||||
print(f" Step {step_count}: cmd=[{command[0]:.2f}, {command[1]:.2f}, {command[2]:.2f}] | "
|
||||
f"trunk_z={trunk_z:.3f}m, vel={lin_vel:.3f}m/s")
|
||||
|
||||
keyboard_reader.restore()
|
||||
view.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,231 +0,0 @@
|
||||
#!/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) # 重置状态 │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
""")
|
||||
@@ -1,25 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Shows how to toss a capsule to a container.
|
||||
"""
|
||||
from re import T
|
||||
from mujoco_py import load_model_from_path, MjSim, MjViewer
|
||||
import os
|
||||
import math
|
||||
import keyboard
|
||||
|
||||
model = load_model_from_path("/path/to/a1.xml")
|
||||
sim = MjSim(model)
|
||||
|
||||
viewer = MjViewer(sim)
|
||||
|
||||
sim_state = sim.get_state()
|
||||
|
||||
while True:
|
||||
|
||||
sim.step()
|
||||
viewer.render()
|
||||
#print(" ")
|
||||
|
||||
if os.getenv('TESTING') is not None:
|
||||
break
|
||||
Reference in New Issue
Block a user