[software] 添加16DOF早期训练仿真与Sim2Real闭环
This commit is contained in:
@@ -0,0 +1,438 @@
|
||||
"""
|
||||
Sim2Sim: Deploy rc_mjlab policy in MuJoCo with RC_MAP terrain.
|
||||
|
||||
Actuator setup matches mjlab training exactly:
|
||||
- Legs: <position> actuator (kp=40, kd=1) — d.ctrl = target_position
|
||||
- Wheels: <velocity> actuator (kd=0.5) — d.ctrl = target_velocity
|
||||
|
||||
The original XML's <general> actuators are overridden in Python to match
|
||||
the mjlab BuiltinPositionActuator / BuiltinVelocityActuator setup.
|
||||
"""
|
||||
import os
|
||||
import time
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import mujoco
|
||||
import mujoco.viewer
|
||||
import numpy as np
|
||||
import pygame
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Policy Model
|
||||
# ============================================================
|
||||
class PolicyMLP(nn.Module):
|
||||
def __init__(self, obs_dim=318, action_dim=16):
|
||||
super().__init__()
|
||||
self.register_buffer("obs_mean", torch.zeros(obs_dim))
|
||||
self.register_buffer("obs_std", torch.ones(obs_dim))
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(obs_dim, 512), nn.ELU(),
|
||||
nn.Linear(512, 256), nn.ELU(),
|
||||
nn.Linear(256, 128), nn.ELU(),
|
||||
nn.Linear(128, action_dim),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = (x - self.obs_mean) / self.obs_std
|
||||
return self.net(x)
|
||||
|
||||
|
||||
def load_policy(model_path, device):
|
||||
ckpt = torch.load(model_path, map_location=device, weights_only=False)
|
||||
state_dict = ckpt["actor_state_dict"]
|
||||
model = PolicyMLP()
|
||||
my_sd = {}
|
||||
for k, v in state_dict.items():
|
||||
if k.startswith("mlp."):
|
||||
my_sd[k.replace("mlp.", "net.")] = v
|
||||
elif k == "obs_normalizer._mean":
|
||||
my_sd["obs_mean"] = v.squeeze()
|
||||
elif k == "obs_normalizer._var":
|
||||
my_sd["obs_std"] = torch.sqrt(v.squeeze() + 1e-5)
|
||||
model.load_state_dict(my_sd, strict=False)
|
||||
model.eval()
|
||||
model.to(device)
|
||||
return model
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Math Utilities
|
||||
# ============================================================
|
||||
def get_gravity_orientation(quat_wxyz):
|
||||
"""
|
||||
Compute projected gravity in body frame from quaternion [w,x,y,z].
|
||||
Proven formula from DreamWaQ-sim2sim reference (easy_math.py).
|
||||
"""
|
||||
qw, qx, qy, qz = quat_wxyz
|
||||
gx = 2.0 * (-qz * qx + qw * qy)
|
||||
gy = -2.0 * (qz * qy + qw * qx)
|
||||
gz = 1.0 - 2.0 * (qw * qw + qz * qz)
|
||||
return np.array([gx, gy, gz], dtype=np.float32)
|
||||
|
||||
|
||||
def quat_rotate_inverse(quat_wxyz, v):
|
||||
"""
|
||||
Rotate vector v from world frame to body frame.
|
||||
quat_wxyz: [w, x, y, z] (as stored in MuJoCo qpos[3:7])
|
||||
v: [3] world-frame vector
|
||||
Same formula as go2w_sim2sim/lab2mujoco.py world2self (no conjugate).
|
||||
"""
|
||||
q_w = quat_wxyz[0]
|
||||
q_vec = quat_wxyz[1:]
|
||||
a = v * (2.0 * q_w * q_w - 1.0)
|
||||
b = np.cross(q_vec, v) * q_w * 2.0
|
||||
c = q_vec * np.dot(q_vec, v) * 2.0
|
||||
return a - b + c
|
||||
|
||||
|
||||
# ============================================================
|
||||
# XML Preparation
|
||||
# ============================================================
|
||||
def create_sim2sim_xml(terrain_xml_path, robot_xml_path, hfield_dir, out_xml_path):
|
||||
with open(terrain_xml_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
|
||||
robot_xml_abs = str(robot_xml_path.absolute()).replace("\\", "/")
|
||||
content = content.replace(
|
||||
'<include file="go2w.xml"/>', f'<include file="{robot_xml_abs}"/>'
|
||||
)
|
||||
|
||||
hfield_1 = str((hfield_dir / "height_field.png").absolute()).replace("\\", "/")
|
||||
hfield_2 = str((hfield_dir / "unitree_hfield.png").absolute()).replace("\\", "/")
|
||||
content = content.replace("../height_field.png", hfield_1)
|
||||
content = content.replace("../unitree_hfield.png", hfield_2)
|
||||
|
||||
with open(out_xml_path, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
|
||||
|
||||
class LowPassFilter:
|
||||
def __init__(self, cutoff_freq, dt, dim):
|
||||
self.alpha = dt / (dt + 1.0 / (2.0 * math.pi * cutoff_freq))
|
||||
self.y_prev = np.zeros(dim, dtype=np.float64)
|
||||
|
||||
def filter(self, x):
|
||||
y = self.alpha * x + (1.0 - self.alpha) * self.y_prev
|
||||
self.y_prev = y.copy()
|
||||
return y
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Main
|
||||
# ============================================================
|
||||
def main():
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
project_root = Path(__file__).parent.parent.absolute()
|
||||
# Local terrain directory to make rc_mjlab independent
|
||||
terrain_dir = Path(__file__).parent / "terrain"
|
||||
terrain_xml = terrain_dir / "scene_terrain.xml"
|
||||
robot_xml = Path(__file__).parent.parent / "mjcf" / "wheelleg.xml"
|
||||
policy_path = Path(__file__).parent.parent / "model_1700.pt"
|
||||
hfield_dir = terrain_dir
|
||||
|
||||
temp_xml = project_root / "mjcf" / "sim2sim_temp.xml"
|
||||
create_sim2sim_xml(terrain_xml, robot_xml, hfield_dir, temp_xml)
|
||||
|
||||
print("Loading MuJoCo model...")
|
||||
os.chdir(str(project_root / "mjcf"))
|
||||
|
||||
# Build model using MjSpec (same as mjlab training) to avoid
|
||||
# broken <include> which renames joints and drops actuators.
|
||||
spec = mujoco.MjSpec.from_file(str(temp_xml))
|
||||
|
||||
# Delete existing XML actuators (mjlab's get_spec() does this)
|
||||
actuators_to_delete = list(spec.actuators)
|
||||
for act in actuators_to_delete:
|
||||
spec.delete(act)
|
||||
|
||||
# Rebuild actuators matching mjlab training config exactly
|
||||
KP_LEG = 40.0
|
||||
KD_LEG = 1.0
|
||||
KD_WHEEL = 0.5
|
||||
EFFORT_LIMIT = 17.0
|
||||
|
||||
leg_joint_names = [
|
||||
"fl_hip_abduction_joint", "fr_hip_abduction_joint",
|
||||
"rl_hip_abduction_joint", "rr_hip_abduction_joint",
|
||||
"fl_hip_pitch_joint", "fr_hip_pitch_joint",
|
||||
"rl_hip_pitch_joint", "rr_hip_pitch_joint",
|
||||
"fl_knee_joint", "fr_knee_joint",
|
||||
"rl_knee_joint", "rr_knee_joint",
|
||||
]
|
||||
wheel_joint_names = [
|
||||
"fl_wheel_joint", "fr_wheel_joint",
|
||||
"rl_wheel_joint", "rr_wheel_joint",
|
||||
]
|
||||
|
||||
# Add position actuators for legs (kp=40, kd=1)
|
||||
for jname in leg_joint_names:
|
||||
act = spec.add_actuator(name=jname, target=jname)
|
||||
act.trntype = mujoco.mjtTrn.mjTRN_JOINT
|
||||
act.dyntype = mujoco.mjtDyn.mjDYN_NONE
|
||||
act.gaintype = mujoco.mjtGain.mjGAIN_FIXED
|
||||
act.biastype = mujoco.mjtBias.mjBIAS_AFFINE
|
||||
act.gainprm[0] = KP_LEG
|
||||
act.biasprm[1] = -KP_LEG
|
||||
act.biasprm[2] = -KD_LEG
|
||||
act.forcelimited = True
|
||||
act.forcerange[:] = [-EFFORT_LIMIT, EFFORT_LIMIT]
|
||||
act.inheritrange = 0.0
|
||||
act.ctrllimited = False
|
||||
|
||||
# Add velocity actuators for wheels (kd=0.5)
|
||||
for jname in wheel_joint_names:
|
||||
act = spec.add_actuator(name=jname, target=jname)
|
||||
act.trntype = mujoco.mjtTrn.mjTRN_JOINT
|
||||
act.dyntype = mujoco.mjtDyn.mjDYN_NONE
|
||||
act.gaintype = mujoco.mjtGain.mjGAIN_FIXED
|
||||
act.biastype = mujoco.mjtBias.mjBIAS_AFFINE
|
||||
act.gainprm[0] = KD_WHEEL
|
||||
act.biasprm[2] = -KD_WHEEL
|
||||
act.forcelimited = True
|
||||
act.forcerange[:] = [-EFFORT_LIMIT, EFFORT_LIMIT]
|
||||
act.inheritrange = 0.0
|
||||
act.ctrllimited = False
|
||||
|
||||
# Add an explicit light over the spawn area so it's guaranteed to be bright
|
||||
l = spec.worldbody.add_light()
|
||||
l.pos[:] = [3.7, -9.0, 4.0]
|
||||
l.dir[:] = [0.0, 0.0, -1.0]
|
||||
l.diffuse[:] = [0.8, 0.8, 0.8]
|
||||
l.specular[:] = [0.3, 0.3, 0.3]
|
||||
|
||||
m = spec.compile()
|
||||
|
||||
# Boost global headlight to ensure no dark corners when camera moves
|
||||
m.vis.headlight.ambient[:] = [0.6, 0.6, 0.6]
|
||||
m.vis.headlight.diffuse[:] = [0.8, 0.8, 0.8]
|
||||
|
||||
d = mujoco.MjData(m)
|
||||
|
||||
# Verify actuators
|
||||
print(f"Model: {m.njnt} joints, {m.nu} actuators")
|
||||
for i in range(m.nu):
|
||||
print(f" actuator[{i}] = {mujoco.mj_id2name(m, mujoco.mjtObj.mjOBJ_ACTUATOR, i)}")
|
||||
|
||||
print(f"Loading Policy from {policy_path}...")
|
||||
policy = load_policy(policy_path, device)
|
||||
print("Policy loaded successfully.")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Joint Configuration — names match the XML exactly (MjSpec preserves them)
|
||||
# ------------------------------------------------------------------
|
||||
leg_joint_names = [
|
||||
"fl_hip_abduction_joint", "fl_hip_pitch_joint", "fl_knee_joint",
|
||||
"fr_hip_abduction_joint", "fr_hip_pitch_joint", "fr_knee_joint",
|
||||
"rl_hip_abduction_joint", "rl_hip_pitch_joint", "rl_knee_joint",
|
||||
"rr_hip_abduction_joint", "rr_hip_pitch_joint", "rr_knee_joint",
|
||||
]
|
||||
wheel_joint_names = [
|
||||
"fl_wheel_joint",
|
||||
"fr_wheel_joint",
|
||||
"rl_wheel_joint",
|
||||
"rr_wheel_joint",
|
||||
]
|
||||
# ------------------------------------------------------------------
|
||||
all_joint_names = leg_joint_names + wheel_joint_names
|
||||
|
||||
qpos_ids = np.array([
|
||||
m.jnt_qposadr[mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_JOINT, n)]
|
||||
for n in all_joint_names
|
||||
])
|
||||
qvel_ids = np.array([
|
||||
m.jnt_dofadr[mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_JOINT, n)]
|
||||
for n in all_joint_names
|
||||
])
|
||||
ctrl_ids = np.array([
|
||||
mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_ACTUATOR, n)
|
||||
for n in all_joint_names
|
||||
])
|
||||
|
||||
# Verify all IDs are valid
|
||||
for i, name in enumerate(all_joint_names):
|
||||
jid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_JOINT, name)
|
||||
assert jid >= 0, f"Joint '{name}' not found!"
|
||||
assert ctrl_ids[i] >= 0, f"Actuator '{name}' not found!"
|
||||
print(f"All {len(all_joint_names)} joints and actuators verified.")
|
||||
|
||||
# Default joint positions
|
||||
default_dof_pos = np.array([
|
||||
0.0, 0.9, -1.8, # FL
|
||||
0.0, 0.9, -1.8, # FR
|
||||
0.0, 0.9, -1.8, # RL
|
||||
0.0, 0.9, -1.8, # RR
|
||||
0.0, 0.0, 0.0, 0.0, # wheel
|
||||
], dtype=np.float64)
|
||||
|
||||
# Action scales
|
||||
HIP_SCALE = 0.125
|
||||
LEG_POS_SCALE = 0.25
|
||||
WHEEL_VEL_SCALE = 5.0
|
||||
action_scale = np.array([
|
||||
HIP_SCALE, LEG_POS_SCALE, LEG_POS_SCALE, # FL
|
||||
HIP_SCALE, LEG_POS_SCALE, LEG_POS_SCALE, # FR
|
||||
HIP_SCALE, LEG_POS_SCALE, LEG_POS_SCALE, # RL
|
||||
HIP_SCALE, LEG_POS_SCALE, LEG_POS_SCALE, # RR
|
||||
WHEEL_VEL_SCALE, WHEEL_VEL_SCALE, WHEEL_VEL_SCALE, WHEEL_VEL_SCALE
|
||||
], dtype=np.float64)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Physics & Control
|
||||
# ------------------------------------------------------------------
|
||||
sim_dt = m.opt.timestep
|
||||
decimation = int(round(0.02 / sim_dt)) # 50Hz control
|
||||
control_dt = sim_dt * decimation
|
||||
|
||||
lpf_legs = LowPassFilter(cutoff_freq=5.0, dt=control_dt, dim=12)
|
||||
lpf_wheels = LowPassFilter(cutoff_freq=15.0, dt=control_dt, dim=4)
|
||||
|
||||
# History buffer
|
||||
history_length = 6
|
||||
obs_dim = 53
|
||||
obs_history = np.zeros((history_length, obs_dim), dtype=np.float32)
|
||||
|
||||
last_actions = np.zeros(16, dtype=np.float32)
|
||||
command = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Pygame UI
|
||||
# ------------------------------------------------------------------
|
||||
pygame.init()
|
||||
screen = pygame.display.set_mode((400, 300))
|
||||
pygame.display.set_caption("Go2W Sim2Sim Control")
|
||||
font = pygame.font.Font(pygame.font.get_default_font(), 24)
|
||||
|
||||
def get_obs():
|
||||
"""Build 53-dim observation matching IsaacLab env_cfgs.py actor_terms order."""
|
||||
quat_wxyz = d.qpos[3:7].copy()
|
||||
# MuJoCo d.qvel[3:6] for free joints is ALREADY in body frame
|
||||
# (unlike cvel which is in world frame). No rotation needed.
|
||||
ang_vel_body = d.qvel[3:6].copy()
|
||||
|
||||
# Body-frame angular velocity * scale
|
||||
base_ang_vel = (ang_vel_body * 0.25).astype(np.float32)
|
||||
|
||||
# Projected gravity
|
||||
projected_gravity = get_gravity_orientation(quat_wxyz)
|
||||
|
||||
# Joint states
|
||||
dof_pos = d.qpos[qpos_ids]
|
||||
dof_vel = d.qvel[qvel_ids]
|
||||
|
||||
joint_pos_rel = (dof_pos[:12] - default_dof_pos[:12]).astype(np.float32)
|
||||
joint_vel_leg = (dof_vel[:12] * 0.05).astype(np.float32)
|
||||
wheel_vel = (dof_vel[12:] * 0.05).astype(np.float32)
|
||||
|
||||
obs = np.concatenate([
|
||||
base_ang_vel, # 3
|
||||
projected_gravity, # 3
|
||||
command, # 3
|
||||
joint_pos_rel, # 12
|
||||
joint_vel_leg, # 12
|
||||
wheel_vel, # 4
|
||||
last_actions, # 16
|
||||
]) # total = 53
|
||||
return obs
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Initialize: set default pose and let robot settle
|
||||
# ------------------------------------------------------------------
|
||||
print("Dropping robot to floor...")
|
||||
d.qpos[:3] = [3.7, -9.0, 0.6]
|
||||
d.qpos[qpos_ids] = default_dof_pos
|
||||
|
||||
# Set ctrl to default targets so the PD controller holds the pose
|
||||
d.ctrl[ctrl_ids[:12]] = default_dof_pos[:12] # leg position targets
|
||||
d.ctrl[ctrl_ids[12:]] = 0.0 # wheel velocity targets = 0
|
||||
for _ in range(500):
|
||||
mujoco.mj_step(m, d)
|
||||
|
||||
# Fill history buffer
|
||||
init_obs = get_obs()
|
||||
for i in range(history_length):
|
||||
obs_history[i] = init_obs
|
||||
|
||||
print("Starting control loop...")
|
||||
|
||||
with mujoco.viewer.launch_passive(m, d) as viewer:
|
||||
while viewer.is_running():
|
||||
step_start = time.time()
|
||||
|
||||
# --- Pygame UI ---
|
||||
for event in pygame.event.get():
|
||||
if event.type == pygame.QUIT:
|
||||
viewer.close()
|
||||
pygame.quit()
|
||||
return
|
||||
|
||||
keys = pygame.key.get_pressed()
|
||||
cmd_vx, cmd_vy, cmd_wz = 0.0, 0.0, 0.0
|
||||
if keys[pygame.K_UP]: cmd_vx = 1.0
|
||||
if keys[pygame.K_DOWN]: cmd_vx = -1.0
|
||||
if keys[pygame.K_LEFT]: cmd_vy = 0.5
|
||||
if keys[pygame.K_RIGHT]: cmd_vy = -0.5
|
||||
if keys[pygame.K_a]: cmd_wz = 1.0
|
||||
if keys[pygame.K_d]: cmd_wz = -1.0
|
||||
command[0] = cmd_vx
|
||||
command[1] = cmd_vy
|
||||
command[2] = cmd_wz
|
||||
|
||||
screen.fill((30, 30, 30))
|
||||
screen.blit(font.render("Go2W Sim2Sim Control", True, (255, 255, 255)), (20, 20))
|
||||
screen.blit(font.render(f"VX (UP/DOWN): {cmd_vx:.1f}", True, (0, 255, 0)), (20, 60))
|
||||
screen.blit(font.render(f"VY (LEFT/RIGHT): {cmd_vy:.1f}", True, (0, 255, 0)), (20, 100))
|
||||
screen.blit(font.render(f"WZ (A/D): {cmd_wz:.1f}", True, (0, 255, 0)), (20, 140))
|
||||
screen.blit(font.render(f"Time: {d.time:.1f}s", True, (200, 200, 200)), (20, 200))
|
||||
pygame.display.flip()
|
||||
|
||||
# --- Policy inference at control frequency ---
|
||||
obs = get_obs()
|
||||
# update history buffer
|
||||
obs_history = np.roll(obs_history, -1, axis=0)
|
||||
obs_history[-1] = obs
|
||||
|
||||
# mjlab observation layout: flatten history per-term, then concatenate
|
||||
term_dims = [3, 3, 3, 12, 12, 4, 16]
|
||||
term_histories = np.split(obs_history, np.cumsum(term_dims)[:-1], axis=1)
|
||||
flat_obs = np.concatenate([h.flatten() for h in term_histories])
|
||||
|
||||
# policy inference
|
||||
pi_input = torch.tensor(flat_obs, dtype=torch.float32, device=device).unsqueeze(0)
|
||||
|
||||
with torch.no_grad():
|
||||
actions = policy(pi_input).squeeze(0).cpu().numpy()
|
||||
|
||||
actions = np.clip(actions, -100.0, 100.0)
|
||||
last_actions[:] = actions
|
||||
|
||||
# Scale and filter
|
||||
scaled = actions * action_scale
|
||||
leg_targets = lpf_legs.filter(scaled[:12])
|
||||
wheel_targets = lpf_wheels.filter(scaled[12:])
|
||||
|
||||
# Apply to MuJoCo ctrl:
|
||||
# Legs: position targets (offset by default pose)
|
||||
d.ctrl[ctrl_ids[:12]] = leg_targets + default_dof_pos[:12]
|
||||
# Wheels: velocity targets
|
||||
d.ctrl[ctrl_ids[12:]] = wheel_targets
|
||||
|
||||
# Step simulation (decimation steps per control step)
|
||||
for _ in range(decimation):
|
||||
mujoco.mj_step(m, d)
|
||||
|
||||
viewer.sync()
|
||||
|
||||
elapsed = time.time() - step_start
|
||||
sleep_time = control_dt - elapsed
|
||||
if sleep_time > 0:
|
||||
time.sleep(sleep_time)
|
||||
|
||||
pygame.quit()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user