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

Init

Browse Source
This commit is contained in:
wty-yy
2025-12-29 16:33:09 +08:00
parent 83952bd549
commit 0c56db7e42
96 changed files with 770085 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

61
deploy/deploy_mujoco/configs/go2.yaml Normal file
View File

@@ -0,0 +1,61 @@
policy_path: "{LEGGED_GYM_ROOT_DIR}/deploy/pre_train/go2/go2_cts_150k.pt"
# xml_path: "{LEGGED_GYM_ROOT_DIR}/resources/robots/go2/flat.xml"
# xml_path: "{LEGGED_GYM_ROOT_DIR}/resources/robots/go2/race_track.xml" # change go2 init pos to pos="-5 2 0.445" in resources/robots/go2/go2.xml
xml_path: "{LEGGED_GYM_ROOT_DIR}/resources/robots/go2/stairs.xml"
# Total simulation time
simulation_duration: 60000000.0
# Simulation time step
simulation_dt: 0.002
# Controller update frequency (meets the requirement of simulation_dt * controll_decimation=0.02; 50Hz)
control_decimation: 10
kps: [20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0]
kds: [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
default_angles: [0.1, 0.8, -1.5, -0.1, 0.8, -1.5,
0.1, 1.0, -1.5, -0.1, 1.0, -1.5]
lin_vel_scale: 2.0
ang_vel_scale: 0.25
dof_pos_scale: 1.0
dof_vel_scale: 0.05
action_scale: 0.25
cmd_scale: [2.0, 2.0, 0.25]
num_actions: 12
num_obs: 45
max_cmd: [2.0, 1.5, 2.5]
cmd_init: [1.0, 0.0, 0.0]
mujoco_joint_names: [
"FL_hip_joint",
"FL_thigh_joint",
"FL_calf_joint",
"FR_hip_joint",
"FR_thigh_joint",
"FR_calf_joint",
"RL_hip_joint",
"RL_thigh_joint",
"RL_calf_joint",
"RR_hip_joint",
"RR_thigh_joint",
"RR_calf_joint"
]
# IsaacGym is same as Mujoco joint names
model_joint_names: [
"FL_hip_joint",
"FL_thigh_joint",
"FL_calf_joint",
"FR_hip_joint",
"FR_thigh_joint",
"FR_calf_joint",
"RL_hip_joint",
"RL_thigh_joint",
"RL_calf_joint",
"RR_hip_joint",
"RR_thigh_joint",
"RR_calf_joint"
]

251
deploy/deploy_mujoco/deploy_go2.py Normal file
View File

@@ -0,0 +1,251 @@
import time
import mujoco.viewer
import mujoco
import numpy as np
from legged_gym import LEGGED_GYM_ROOT_DIR
import torch
import yaml
import os
import imageio
from pathlib import Path
from argparse import ArgumentParser
import pygame
from matplotlib import pyplot as plt
def get_gravity_orientation(quaternion):
qw = quaternion[0]
qx = quaternion[1]
qy = quaternion[2]
qz = quaternion[3]
gravity_orientation = np.zeros(3)
gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
gravity_orientation[1] = -2 * (qz * qy + qw * qx)
gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
return gravity_orientation
def pd_control(target_q, q, kp, target_dq, dq, kd):
"""Calculates torques from position commands"""
return (target_q - q) * kp + (target_dq - dq) * kd
def get_xbox_command(joystick, max_cmd):
pygame.event.pump()
dead_zone = 0.1
lx = joystick.get_axis(0)
ly = joystick.get_axis(1)
rx = joystick.get_axis(3)
if abs(lx) < dead_zone: lx = 0
if abs(ly) < dead_zone: ly = 0
if abs(rx) < dead_zone: rx = 0
cmd_x = -ly * max_cmd[0]
cmd_y = -lx * max_cmd[1]
cmd_yaw = -rx * max_cmd[2]
return np.array([cmd_x, cmd_y, cmd_yaw], dtype=np.float32)
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("--save-video", action="store_true", help="Whether to save video of the simulation.")
parser.add_argument("--visualize-moe-weights", action="store_true", help="Whether to visualize mixture of experts weights.")
parser.add_argument("--save-moe-latent", action="store_true", help="Whether to save mixture of experts latent vectors.")
args = parser.parse_args()
save_video = args.save_video
visualize_moe_weights = args.visualize_moe_weights
save_moe_latent = args.save_moe_latent
config_file = "go2.yaml"
pygame.init()
use_joystick = False
joystick = None
if pygame.joystick.get_count() > 0:
joystick = pygame.joystick.Joystick(0)
joystick.init()
use_joystick = True
print(f"Detected Joystick: {joystick.get_name()}")
else:
print("No Joystick detected. Using default commands from config.")
with open(f"{LEGGED_GYM_ROOT_DIR}/deploy/deploy_mujoco/configs/{config_file}", "r") as f:
config = yaml.load(f, Loader=yaml.FullLoader)
policy_path = config["policy_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
xml_path = config["xml_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
simulation_duration = config["simulation_duration"]
simulation_dt = config["simulation_dt"]
control_decimation = config["control_decimation"]
kps = np.array(config["kps"], dtype=np.float32)
kds = np.array(config["kds"], dtype=np.float32)
default_angles = np.array(config["default_angles"], dtype=np.float32)
lin_vel_scale = config["lin_vel_scale"]
ang_vel_scale = config["ang_vel_scale"]
dof_pos_scale = config["dof_pos_scale"]
dof_vel_scale = config["dof_vel_scale"]
action_scale = config["action_scale"]
cmd_scale = np.array(config["cmd_scale"], dtype=np.float32)
num_actions = config["num_actions"]
num_obs = config["num_obs"]
cmd = np.array(config["cmd_init"], dtype=np.float32)
idx_model2mj = idx_mj2model = list(range(num_actions))
if 'mujoco_joint_names' in config and 'model_joint_names' in config:
mujoco_joint_names = config["mujoco_joint_names"]
model_joint_names = config["model_joint_names"]
idx_model2mj = [model_joint_names.index(joint) for joint in mujoco_joint_names]
idx_mj2model = [mujoco_joint_names.index(joint) for joint in model_joint_names]
video_save_dir = str(Path(__file__).parent / "videos")
os.makedirs(video_save_dir, exist_ok=True)
model_name = os.path.basename(policy_path).split('.')[0]
cmd_str = f"cmd_{cmd[0]}_{cmd[1]}_{cmd[2]}"
# define context variables
action = np.zeros(num_actions, dtype=np.float32)
last_action = np.zeros(num_actions, dtype=np.float32)
target_dof_pos = default_angles.copy()
obs = np.zeros(num_obs, dtype=np.float32)
counter = 0
# Load robot model
m = mujoco.MjModel.from_xml_path(xml_path)
d = mujoco.MjData(m)
m.opt.timestep = simulation_dt
renderer = mujoco.Renderer(m, height=360, width=640)
# load policy
policy = torch.jit.load(policy_path)
if save_video:
video_filename = f"{model_name}_{cmd_str}.mp4"
video_path = os.path.join(video_save_dir, video_filename)
print(f"Video recording will be saved to: {video_path}")
video_fps = 50
sim_fps = 1.0 / m.opt.timestep
frame_skip = int(sim_fps / video_fps)
if frame_skip < 1:
frame_skip = 1
writer = imageio.get_writer(video_path, fps=video_fps)
print(f"Sim FPS: {sim_fps:.2f}, Video FPS: {video_fps}, Frame Skip: {frame_skip}, Save at: {video_path}")
if visualize_moe_weights:
plt.ion()
fig, ax = plt.subplots(figsize=(5,3))
ax.set_title(f"Command: Vx={cmd[0]:.2f}, Vy={cmd[1]:.2f}, Wz={cmd[2]:.2f}")
bars = None
if save_moe_latent:
latent_save_dir = str(Path(__file__).parent / "data_latents")
os.makedirs(latent_save_dir, exist_ok=True)
latent_filename = f"{model_name}_{cmd_str}_latents.npy"
latent_path = os.path.join(latent_save_dir, latent_filename)
all_latents = []
with mujoco.viewer.launch_passive(m, d) as viewer:
# set viewer.camera to follow robot
viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
viewer.cam.trackbodyid = 1
viewer.cam.distance = 3.0
viewer.cam.elevation = -30.0
viewer.cam.azimuth = 0.0
# Close the viewer automatically after simulation_duration wall-seconds.
start = time.time()
while viewer.is_running() and time.time() - start < simulation_duration:
step_start = time.time()
if use_joystick and counter % control_decimation == 0:
cmd = get_xbox_command(joystick, config["max_cmd"])
print(f"Cmd: Vx={cmd[0]:.2f}, Vy={cmd[1]:.2f}, Wz={cmd[2]:.2f}", end='\r')
tau = pd_control(target_dof_pos, d.qpos[7:], kps, np.zeros_like(kds), d.qvel[6:], kds)
d.ctrl[:] = tau
# mj_step can be replaced with code that also evaluates
# a policy and applies a control signal before stepping the physics.
mujoco.mj_step(m, d)
if save_video and counter % frame_skip == 0:
try:
renderer.update_scene(d, camera=viewer.cam)
frame = renderer.render()
writer.append_data(frame)
except Exception as e:
print(f"Error rendering frame: {e}")
counter += 1
if counter % control_decimation == 0:
# Apply control signal here.
# create observation
qj = d.qpos[7:]
dqj = d.qvel[6:]
quat = d.qpos[3:7]
lin_vel = d.qvel[:3]
ang_vel = d.qvel[3:6]
qj = (qj - default_angles) * dof_pos_scale
dqj = dqj * dof_vel_scale
gravity_orientation = get_gravity_orientation(quat)
lin_vel = lin_vel * lin_vel_scale
ang_vel = ang_vel * ang_vel_scale
obs[:3] = ang_vel
obs[3:6] = gravity_orientation
obs[6:9] = cmd * cmd_scale
obs[9 : 9 + num_actions] = qj[idx_mj2model]
obs[9 + num_actions : 9 + 2 * num_actions] = dqj[idx_mj2model]
obs[9 + 2 * num_actions : 9 + 3 * num_actions] = action[idx_mj2model]
obs_tensor = torch.from_numpy(obs).unsqueeze(0)
# policy inference
last_action = action
result = policy(obs_tensor)
if isinstance(result, tuple):
action, (weights, latent) = result # moe
action = action.detach().numpy().squeeze()[idx_model2mj]
weights = weights.detach().numpy().squeeze()
latent = latent.detach().numpy().squeeze()
if visualize_moe_weights:
if bars is None:
x = np.arange(len(weights))
bars = ax.bar(x, weights)
ax.set_ylim(0, 1)
else:
for bar, w in zip(bars, weights):
bar.set_height(w)
plt.draw()
plt.pause(0.001) # 这会造成大约 1ms 的延迟
if save_moe_latent:
all_latents.append(latent)
else:
action = result.detach().cpu().numpy().squeeze()[idx_model2mj]
# transform action to target_dof_pos
target_dof_pos = action * action_scale + default_angles
# Pick up changes to the physics state, apply perturbations, update options from GUI.
viewer.sync()
# Rudimentary time keeping, will drift relative to wall clock.
# time_until_next_step = m.opt.timestep - (time.time() - step_start) - 0.1
# if time_until_next_step > 0:
# time.sleep(time_until_next_step)
# writer.close()
if save_video:
print(f"Video saved successfully to {video_path}")
writer.close()
if save_moe_latent and len(all_latents) > 0:
all_latents = np.array(all_latents)
np.save(latent_path, all_latents)
print(f"Latent vectors saved successfully to {latent_path}")

126
deploy/deploy_mujoco/deploy_go2_lab.py Normal file
View File

@@ -0,0 +1,126 @@
import time
import mujoco.viewer
import mujoco
import numpy as np
from legged_gym import LEGGED_GYM_ROOT_DIR
import torch
import yaml
def get_gravity_orientation(quaternion):
qw = quaternion[0]
qx = quaternion[1]
qy = quaternion[2]
qz = quaternion[3]
gravity_orientation = np.zeros(3)
gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
gravity_orientation[1] = -2 * (qz * qy + qw * qx)
gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
return gravity_orientation
def pd_control(target_q, q, kp, target_dq, dq, kd):
"""Calculates torques from position commands"""
return (target_q - q) * kp + (target_dq - dq) * kd
if __name__ == "__main__":
# get config file name from command line
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("config_file", type=str, help="config file name in the config folder")
args = parser.parse_args()
config_file = args.config_file
with open(f"{LEGGED_GYM_ROOT_DIR}/deploy/deploy_mujoco/configs/{config_file}", "r") as f:
config = yaml.load(f, Loader=yaml.FullLoader)
policy_path = config["policy_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
xml_path = config["xml_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
simulation_duration = config["simulation_duration"]
simulation_dt = config["simulation_dt"]
control_decimation = config["control_decimation"]
kps = np.array(config["kps"], dtype=np.float32)
kds = np.array(config["kds"], dtype=np.float32)
default_angles = np.array(config["default_angles"], dtype=np.float32)
joint_ids_map = config["joint_ids_map"]
ang_vel_scale = config["ang_vel_scale"]
dof_pos_scale = config["dof_pos_scale"]
dof_vel_scale = config["dof_vel_scale"]
action_scale = config["action_scale"]
cmd_scale = np.array(config["cmd_scale"], dtype=np.float32)
num_actions = config["num_actions"]
num_obs = config["num_obs"]
cmd = np.array(config["cmd_init"], dtype=np.float32)
# define context variables
action = np.zeros(num_actions, dtype=np.float32)
target_dof_pos = default_angles.copy()
obs = np.zeros(num_obs, dtype=np.float32)
counter = 0
# Load robot model
m = mujoco.MjModel.from_xml_path(xml_path)
d = mujoco.MjData(m)
m.opt.timestep = simulation_dt
# load policy
policy = torch.jit.load(policy_path)
with mujoco.viewer.launch_passive(m, d) as viewer:
# Close the viewer automatically after simulation_duration wall-seconds.
start = time.time()
while viewer.is_running() and time.time() - start < simulation_duration:
step_start = time.time()
temp = target_dof_pos[[0,4,8,1,5,9,2,6,10,3,7,11]]
tau = pd_control(temp, d.qpos[7:], kps, np.zeros_like(kds), d.qvel[6:], kds)
d.ctrl[:] = tau
# mj_step can be replaced with code that also evaluates
# a policy and applies a control signal before stepping the physics.
mujoco.mj_step(m, d)
counter += 1
if counter % control_decimation == 0:
# Apply control signal here.
# create observation
qj = d.qpos[7:]
dqj = d.qvel[6:]
quat = d.qpos[3:7]
ang_vel = d.qvel[3:6]
qj = (qj - default_angles) * dof_pos_scale
dqj = dqj * dof_vel_scale
gravity_orientation = get_gravity_orientation(quat)
ang_vel = ang_vel * ang_vel_scale
obs[:3] = ang_vel
obs[3:6] = gravity_orientation
obs[6:9] = cmd * cmd_scale
obs[9 : 9 + num_actions] = qj[joint_ids_map]
obs[9 + num_actions : 9 + 2 * num_actions] = dqj[joint_ids_map]
obs[9 + 2 * num_actions : 9 + 3 * num_actions] = action
obs_tensor = torch.from_numpy(obs).unsqueeze(0)
# policy inference
action = policy(obs_tensor).detach().numpy().squeeze()
# transform action to target_dof_pos
target_dof_pos = action * action_scale + default_angles[joint_ids_map]
# Pick up changes to the physics state, apply perturbations, update options from GUI.
viewer.sync()
# Rudimentary time keeping, will drift relative to wall clock.
time_until_next_step = m.opt.timestep - (time.time() - step_start)
if time_until_next_step > 0:
time.sleep(time_until_next_step)

272
deploy/deploy_mujoco/deploy_go2_moe.py Normal file
View File

@@ -0,0 +1,272 @@
import time
import mujoco.viewer
import mujoco
import numpy as np
from legged_gym import LEGGED_GYM_ROOT_DIR
import torch
import yaml
import os
import imageio
from pathlib import Path
from argparse import ArgumentParser
import pygame
# from matplotlib import pyplot as plt # 移除 matplotlib
def get_gravity_orientation(quaternion):
qw = quaternion[0]
qx = quaternion[1]
qy = quaternion[2]
qz = quaternion[3]
gravity_orientation = np.zeros(3)
gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
gravity_orientation[1] = -2 * (qz * qy + qw * qx)
gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
return gravity_orientation
def pd_control(target_q, q, kp, target_dq, dq, kd):
"""Calculates torques from position commands"""
return (target_q - q) * kp + (target_dq - dq) * kd
def get_xbox_command(joystick, max_cmd):
# 注意：如果开启了 Pygame 显示窗口，这里 event.pump 也是必要的
pygame.event.pump()
dead_zone = 0.1
if joystick is not None:
lx = joystick.get_axis(0)
ly = joystick.get_axis(1)
rx = joystick.get_axis(3)
if abs(lx) < dead_zone: lx = 0
if abs(ly) < dead_zone: ly = 0
if abs(rx) < dead_zone: rx = 0
cmd_x = -ly * max_cmd[0]
cmd_y = -lx * max_cmd[1]
cmd_yaw = -rx * max_cmd[2]
return np.array([cmd_x, cmd_y, cmd_yaw], dtype=np.float32)
return np.zeros(3, dtype=np.float32)
def draw_moe_weights(screen, weights, width, height):
"""使用 Pygame 绘制 MoE 权重"""
screen.fill((255, 255, 255)) # 白底
num_experts = len(weights)
if num_experts == 0:
return
# 设置边距
margin = 5
bar_width = (width - 2 * margin) / num_experts
max_bar_height = height - 2 * margin
for i, w in enumerate(weights):
# 限制 w 在 [0, 1] 之间用于显示
w_clamped = max(0.0, min(1.0, w))
bar_height = int(w_clamped * max_bar_height)
# 计算矩形位置 (Pygame 坐标原点在左上角)
# left, top, width, height
x = margin + i * bar_width
y = height - margin - bar_height # 从底部向上长
# 绘制矩形 (蓝色)
# 在 bar 之间留一点空隙 (width - 2)
pygame.draw.rect(screen, (50, 100, 255), (x, y, bar_width - 2, bar_height))
pygame.display.flip()
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("--save-video", action="store_true", help="Whether to save video of the simulation.")
parser.add_argument("--visualize-moe-weights", action="store_true", help="Whether to visualize mixture of experts weights.")
args = parser.parse_args()
save_video = args.save_video
visualize_moe_weights = args.visualize_moe_weights
config_file = "go2.yaml"
# Pygame 初始化
pygame.init()
use_joystick = False
joystick = None
if pygame.joystick.get_count() > 0:
joystick = pygame.joystick.Joystick(0)
joystick.init()
use_joystick = True
print(f"Detected Joystick: {joystick.get_name()}")
else:
print("No Joystick detected. Using default commands from config.")
# 如果需要可视化权重，设置 Pygame 窗口
screen = None
win_width, win_height = 400, 200
if visualize_moe_weights:
# 创建一个独立的窗口用于显示权重
screen = pygame.display.set_mode((win_width, win_height))
pygame.display.set_caption("MoE Weights Visualization")
with open(f"{LEGGED_GYM_ROOT_DIR}/deploy/deploy_mujoco/configs/{config_file}", "r") as f:
config = yaml.load(f, Loader=yaml.FullLoader)
policy_path = config["policy_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
xml_path = config["xml_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
simulation_duration = config["simulation_duration"]
simulation_dt = config["simulation_dt"]
control_decimation = config["control_decimation"]
kps = np.array(config["kps"], dtype=np.float32)
kds = np.array(config["kds"], dtype=np.float32)
default_angles = np.array(config["default_angles"], dtype=np.float32)
lin_vel_scale = config["lin_vel_scale"]
ang_vel_scale = config["ang_vel_scale"]
dof_pos_scale = config["dof_pos_scale"]
dof_vel_scale = config["dof_vel_scale"]
action_scale = config["action_scale"]
cmd_scale = np.array(config["cmd_scale"], dtype=np.float32)
num_actions = config["num_actions"]
num_obs = config["num_obs"]
cmd = np.array(config["cmd_init"], dtype=np.float32)
idx_model2mj = idx_mj2model = list(range(num_actions))
if 'mujoco_joint_names' in config and 'model_joint_names' in config:
mujoco_joint_names = config["mujoco_joint_names"]
model_joint_names = config["model_joint_names"]
idx_model2mj = [model_joint_names.index(joint) for joint in mujoco_joint_names]
idx_mj2model = [mujoco_joint_names.index(joint) for joint in model_joint_names]
video_save_dir = str(Path(__file__).parent / "videos")
os.makedirs(video_save_dir, exist_ok=True)
model_name = os.path.basename(policy_path).split('.')[0]
cmd_str = f"cmd_{cmd[0]}_{cmd[1]}_{cmd[2]}"
video_filename = f"{model_name}_{cmd_str}.mp4"
video_path = os.path.join(video_save_dir, video_filename)
print(f"Video recording will be saved to: {video_path}")
# define context variables
action = np.zeros(num_actions, dtype=np.float32)
last_action = np.zeros(num_actions, dtype=np.float32)
target_dof_pos = default_angles.copy()
obs = np.zeros(num_obs, dtype=np.float32)
counter = 0
# Load robot model
m = mujoco.MjModel.from_xml_path(xml_path)
d = mujoco.MjData(m)
m.opt.timestep = simulation_dt
renderer = mujoco.Renderer(m, height=360, width=640)
# load policy
policy = torch.jit.load(policy_path)
if save_video:
video_fps = 50
sim_fps = 1.0 / m.opt.timestep
frame_skip = int(sim_fps / video_fps)
if frame_skip < 1:
frame_skip = 1
writer = imageio.get_writer(video_path, fps=video_fps)
print(f"Sim FPS: {sim_fps:.2f}, Video FPS: {video_fps}, Frame Skip: {frame_skip}, Save at: {video_path}")
# 移除了 plt 初始化逻辑
with mujoco.viewer.launch_passive(m, d) as viewer:
# set viewer.camera to follow robot
viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
viewer.cam.trackbodyid = 1
viewer.cam.distance = 3.0
viewer.cam.elevation = -30.0
viewer.cam.azimuth = 0.0
# Close the viewer automatically after simulation_duration wall-seconds.
start = time.time()
while viewer.is_running() and time.time() - start < simulation_duration:
step_start = time.time()
if use_joystick and counter % control_decimation == 0:
cmd = get_xbox_command(joystick, config["max_cmd"])
print(f"Cmd: Vx={cmd[0]:.2f}, Vy={cmd[1]:.2f}, Wz={cmd[2]:.2f}", end='\r')
elif visualize_moe_weights and counter % control_decimation == 0:
# 如果没有手柄但开了可视化，也需要 pump 事件，防止窗口卡死
pygame.event.pump()
tau = pd_control(target_dof_pos, d.qpos[7:], kps, np.zeros_like(kds), d.qvel[6:], kds)
d.ctrl[:] = tau
mujoco.mj_step(m, d)
if save_video and counter % frame_skip == 0:
try:
renderer.update_scene(d, camera=viewer.cam)
frame = renderer.render()
writer.append_data(frame)
except Exception as e:
print(f"Error rendering frame: {e}")
counter += 1
if counter % control_decimation == 0:
# Apply control signal here.
# create observation
qj = d.qpos[7:]
dqj = d.qvel[6:]
quat = d.qpos[3:7]
lin_vel = d.qvel[:3]
ang_vel = d.qvel[3:6]
qj = (qj - default_angles) * dof_pos_scale
dqj = dqj * dof_vel_scale
gravity_orientation = get_gravity_orientation(quat)
lin_vel = lin_vel * lin_vel_scale
ang_vel = ang_vel * ang_vel_scale
obs[:3] = ang_vel
obs[3:6] = gravity_orientation
obs[6:9] = cmd * cmd_scale
obs[9 : 9 + num_actions] = qj[idx_mj2model]
obs[9 + num_actions : 9 + 2 * num_actions] = dqj[idx_mj2model]
obs[9 + 2 * num_actions : 9 + 3 * num_actions] = action[idx_mj2model]
obs_tensor = torch.from_numpy(obs).unsqueeze(0)
# policy inference
last_action = action
result = policy(obs_tensor)
# 处理 MoE 和 绘图
if isinstance(result, tuple):
action, weights = result # moe
action = action.detach().numpy().squeeze()[idx_model2mj]
weights = weights.detach().numpy().squeeze()
if visualize_moe_weights and screen is not None:
draw_moe_weights(screen, weights, win_width, win_height)
else:
action = result.detach().numpy().squeeze()[idx_model2mj]
# transform action to target_dof_pos
target_dof_pos = action * action_scale + default_angles
# Pick up changes to the physics state, apply perturbations, update options from GUI.
viewer.sync()
# 如果需要严格同步时间，可以解开下面的注释
# time_until_next_step = m.opt.timestep - (time.time() - step_start)
# if time_until_next_step > 0:
# time.sleep(time_until_next_step)
if save_video:
writer.close()
# 退出时清理 Pygame
pygame.quit()
print(f"Video saved successfully to {video_path}")

61
deploy/deploy_real/common/command_helper.py Normal file
View File

@@ -0,0 +1,61 @@
from unitree_sdk2py.idl.unitree_go.msg.dds_ import LowCmd_ as LowCmdGo
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_ as LowCmdHG
from typing import Union
class MotorMode:
PR = 0 # Series Control for Pitch/Roll Joints
AB = 1 # Parallel Control for A/B Joints
def create_damping_cmd(cmd: Union[LowCmdGo, LowCmdHG]):
size = len(cmd.motor_cmd)
for i in range(size):
cmd.motor_cmd[i].q = 0
cmd.motor_cmd[i].qd = 0
cmd.motor_cmd[i].kp = 0
cmd.motor_cmd[i].kd = 8
cmd.motor_cmd[i].tau = 0
def create_zero_cmd(cmd: Union[LowCmdGo, LowCmdHG]):
size = len(cmd.motor_cmd)
for i in range(size):
cmd.motor_cmd[i].q = 0
cmd.motor_cmd[i].qd = 0
cmd.motor_cmd[i].kp = 0
cmd.motor_cmd[i].kd = 0
cmd.motor_cmd[i].tau = 0
def init_cmd_hg(cmd: LowCmdHG, mode_machine: int, mode_pr: int):
cmd.mode_machine = mode_machine
cmd.mode_pr = mode_pr
size = len(cmd.motor_cmd)
for i in range(size):
cmd.motor_cmd[i].mode = 1
cmd.motor_cmd[i].q = 0
cmd.motor_cmd[i].qd = 0
cmd.motor_cmd[i].kp = 0
cmd.motor_cmd[i].kd = 0
cmd.motor_cmd[i].tau = 0
def init_cmd_go(cmd: LowCmdGo, weak_motor: list):
cmd.head[0] = 0xFE
cmd.head[1] = 0xEF
cmd.level_flag = 0xFF
cmd.gpio = 0
PosStopF = 2.146e9
VelStopF = 16000.0
size = len(cmd.motor_cmd)
for i in range(size):
if i in weak_motor:
cmd.motor_cmd[i].mode = 1
else:
cmd.motor_cmd[i].mode = 0x0A
cmd.motor_cmd[i].q = PosStopF
cmd.motor_cmd[i].qd = VelStopF
cmd.motor_cmd[i].kp = 0
cmd.motor_cmd[i].kd = 0
cmd.motor_cmd[i].tau = 0

39
deploy/deploy_real/common/remote_controller.py Normal file
View File

@@ -0,0 +1,39 @@
import struct
class KeyMap:
R1 = 0
L1 = 1
start = 2
select = 3
R2 = 4
L2 = 5
F1 = 6
F2 = 7
A = 8
B = 9
X = 10
Y = 11
up = 12
right = 13
down = 14
left = 15
class RemoteController:
def __init__(self):
self.lx = 0
self.ly = 0
self.rx = 0
self.ry = 0
self.button = [0] * 16
def set(self, data):
# wireless_remote
keys = struct.unpack("H", data[2:4])[0]
for i in range(16):
self.button[i] = (keys & (1 << i)) >> i
self.lx = struct.unpack("f", data[4:8])[0]
self.rx = struct.unpack("f", data[8:12])[0]
self.ry = struct.unpack("f", data[12:16])[0]
self.ly = struct.unpack("f", data[20:24])[0]

25
deploy/deploy_real/common/rotation_helper.py Normal file
View File

@@ -0,0 +1,25 @@
import numpy as np
from scipy.spatial.transform import Rotation as R
def get_gravity_orientation(quaternion):
qw = quaternion[0]
qx = quaternion[1]
qy = quaternion[2]
qz = quaternion[3]
gravity_orientation = np.zeros(3)
gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
gravity_orientation[1] = -2 * (qz * qy + qw * qx)
gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
return gravity_orientation
def transform_imu_data(waist_yaw, waist_yaw_omega, imu_quat, imu_omega):
RzWaist = R.from_euler("z", waist_yaw).as_matrix()
R_torso = R.from_quat([imu_quat[1], imu_quat[2], imu_quat[3], imu_quat[0]]).as_matrix()
R_pelvis = np.dot(R_torso, RzWaist.T)
w = np.dot(RzWaist, imu_omega[0]) - np.array([0, 0, waist_yaw_omega])
return R.from_matrix(R_pelvis).as_quat()[[3, 0, 1, 2]], w

35
deploy/deploy_real/config_go2.py Normal file
View File

@@ -0,0 +1,35 @@
from legged_gym import LEGGED_GYM_ROOT_DIR
import numpy as np
import yaml
class Config:
def __init__(self, file_path) -> None:
with open(file_path, "r") as f:
config = yaml.load(f, Loader=yaml.FullLoader)
self.control_dt = config["control_dt"]
self.joint2motor_idx = config["joint2motor_idx"]
self.msg_type = config["msg_type"]
self.imu_type = config["imu_type"]
self.lowcmd_topic = config["lowcmd_topic"]
self.lowstate_topic = config["lowstate_topic"]
self.policy_path = config["policy_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
self.kps = np.array(config["kps"],dtype=np.float32)
self.kds = np.array(config["kds"],dtype=np.float32)
self.default_angles = np.array(config["default_angles"], dtype=np.float32)
self.obs_scales_ang_vel = config["obs_scales_ang_vel"]
self.obs_scales_dof_pos = config["obs_scales_dof_pos"]
self.obs_scales_dof_vel = config["obs_scales_dof_vel"]
self.command_scale = config["command_scale"]
self.action_scale = config["action_scale"]
self.num_actions = config["num_actions"]
self.num_obs = config["num_obs"]

27
deploy/deploy_real/configs/go2.yaml Normal file
View File

@@ -0,0 +1,27 @@
control_dt: 0.02
msg_type: "go" # "hg" or "go"
imu_type: "torso" # "torso" or "pelvis"
lowcmd_topic: "rt/lowcmd"
lowstate_topic: "rt/lowstate"
policy_path: "{LEGGED_GYM_ROOT_DIR}/deploy/pre_train/go2/go2_cts_150k.pt"
joint2motor_idx: [3,4,5,0,1,2,9,10,11,6,7,8]
kps: [20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20]
kds: [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
default_angles: [ 0.1, 0.8, -1.5,
-0.1, 0.8, -1.5,
0.1, 1.0, -1.5,
-0.1, 1.0, -1.5]
obs_scales_ang_vel: 0.25
obs_scales_dof_pos: 1.0
obs_scales_dof_vel: 0.05
command_scale: [3.0, 2.0, 0.5]
action_scale: 0.25
num_actions: 12
num_obs: 45

218
deploy/deploy_real/deploy_real_go2.py Normal file
View File

@@ -0,0 +1,218 @@
from legged_gym import LEGGED_GYM_ROOT_DIR
import numpy as np
import time
import torch
from unitree_sdk2py.core.channel import ChannelPublisher,ChannelSubscriber,ChannelFactoryInitialize
from unitree_sdk2py.idl.default import unitree_go_msg_dds__LowCmd_,unitree_go_msg_dds__LowState_
from unitree_sdk2py.idl.unitree_go.msg.dds_ import LowCmd_ as LowCmdGo
from unitree_sdk2py.idl.unitree_go.msg.dds_ import LowState_ as LowStateGo
from unitree_sdk2py.utils.crc import CRC
from common.command_helper import create_zero_cmd,create_damping_cmd
from common.rotation_helper import get_gravity_orientation
from common.remote_controller import RemoteController, KeyMap
from config_go2 import Config
HIGHLEVEL = 0xEE
LOWLEVEL = 0xFF
TRIGERLEVEL = 0xF0
PosStopF = 2.146e9
VelStopF = 16000.0
def init_cmd_go2(cmd:LowCmdGo):
cmd.head[0] = 0xFE
cmd.head[1] = 0xEF
cmd.level_flag = 0xFF
cmd.gpio = 0
for i in range(12):
cmd.motor_cmd[i].mode = 0x0A # 0x01
cmd.motor_cmd[i].q = PosStopF
cmd.motor_cmd[i].dq = VelStopF # or qd
cmd.motor_cmd[i].kp = 0.0
cmd.motor_cmd[i].kd = 0.0
cmd.motor_cmd[i].tau = 0.0
class Controller:
def __init__(self,config:Config) -> None:
self.config = config
self.remote_controller = RemoteController()
self.use_remote_controller = True
self.policy = torch.jit.load(config.policy_path)
self._warm_up()
self.qj = np.zeros(config.num_actions,dtype=np.float32)
self.dqj = np.zeros(config.num_actions,dtype=np.float32)
self.action = np.zeros(config.num_actions,dtype=np.float32)
self.target_dof_pos = config.default_angles.copy()
self.obs = np.zeros(config.num_obs,dtype=np.float32)
self.cmd = np.array([0.8, 0, 0],dtype=np.float32)
self.counter = 0
self.low_cmd = unitree_go_msg_dds__LowCmd_()
self.low_state = unitree_go_msg_dds__LowState_()
self.lowcmd_publisher = ChannelPublisher(config.lowcmd_topic,LowCmdGo)
self.lowcmd_publisher.Init()
self.lowstate_subscriber = ChannelSubscriber(config.lowstate_topic,LowStateGo)
self.lowstate_subscriber.Init(self.LowStateHandler,10)
# self.replay_buffer = ReplayBuffer(max_replay_buffer_size=200,flag='real_new')
self.wait_for_low_state()
init_cmd_go2(self.low_cmd)
def _warm_up(self):
obs = torch.ones((1,45))
for _ in range(10):
_ = self.policy(obs)
print('Network has been warmed up.')
def wait_for_low_state(self):
while self.low_state.tick == 0:
time.sleep(self.config.control_dt)
print("Successfully connected to the robot.")
def LowStateHandler(self,msg:LowStateGo):
self.low_state = msg
self.remote_controller.set(self.low_state.wireless_remote)
def send_cmd(self,cmd:LowCmdGo):
cmd.crc = CRC().Crc(cmd)
self.lowcmd_publisher.Write(cmd)
def zero_torque_state(self):
print("Enter zero torque state.")
print("Waiting for the start signal...")
while self.remote_controller.button[KeyMap.start] != 1:
create_zero_cmd(self.low_cmd)
self.send_cmd(self.low_cmd)
time.sleep(self.config.control_dt)
def move_to_default_pos(self):
print('Moving to default pos.')
total_time = 2
num_step = int(total_time / self.config.control_dt)
dof_idx = self.config.joint2motor_idx
default_pos = self.config.default_angles
init_dof_pos = np.zeros(12,dtype=np.float32)
for i in range(12):
init_dof_pos[i] = self.low_state.motor_state[dof_idx[i]].q
for i in range(num_step):
alpha = i / num_step
for j in range(12):
motor_idx = dof_idx[j]
target_pos = default_pos[j]
self.low_cmd.motor_cmd[motor_idx].q = init_dof_pos[j] * (1 - alpha) + target_pos * alpha
self.low_cmd.motor_cmd[motor_idx].dq = 0.0 # qd
self.low_cmd.motor_cmd[motor_idx].kp = 40.0
self.low_cmd.motor_cmd[motor_idx].kd = 0.6
self.low_cmd.motor_cmd[motor_idx].tau = 0.0
self.send_cmd(self.low_cmd)
time.sleep(self.config.control_dt)
def default_pos_state(self):
print("Enter default pos state.")
print("Waiting for the Button A signal...")
while self.remote_controller.button[KeyMap.A] != 1:
for i in range(12):
motor_idx = self.config.joint2motor_idx[i]
self.low_cmd.motor_cmd[motor_idx].q = self.config.default_angles[i]
self.low_cmd.motor_cmd[motor_idx].dq = 0.0 # qd
self.low_cmd.motor_cmd[motor_idx].kp = 40.0
self.low_cmd.motor_cmd[motor_idx].kd = 0.6
self.low_cmd.motor_cmd[motor_idx].tau = 0
self.send_cmd(self.low_cmd)
time.sleep(self.config.control_dt)
def run(self):
self.counter += 1
for i in range(12):
self.qj[i] = self.low_state.motor_state[self.config.joint2motor_idx[i]].q
self.dqj[i] = self.low_state.motor_state[self.config.joint2motor_idx[i]].dq
ang_vel = np.array([self.low_state.imu_state.gyroscope], dtype=np.float32) * self.config.obs_scales_ang_vel
quat = self.low_state.imu_state.quaternion
gravity_orientation = get_gravity_orientation(quat) # imu_state quaternion: w, x, y, z
if self.use_remote_controller:
self.cmd[0] = self.remote_controller.ly
self.cmd[1] = self.remote_controller.lx * -1
self.cmd[2] = self.remote_controller.rx * -1
qj_obs = self.qj.copy()
qj_obs = (qj_obs - self.config.default_angles) * self.config.obs_scales_dof_pos
dqj_obs = self.dqj.copy()
dqj_obs = dqj_obs * self.config.obs_scales_dof_vel
self.obs[:3] = ang_vel
self.obs[3:6] = gravity_orientation
self.obs[6:9] = self.cmd * self.config.command_scale
self.obs[9:21] = qj_obs
self.obs[21:33] = dqj_obs
self.obs[33:45] = self.action
obs_tensor = torch.from_numpy(self.obs).unsqueeze(0)
self.action = self.policy(obs_tensor).detach().numpy().squeeze()
target_dof_pos = self.config.default_angles + self.action * self.config.action_scale
# target_dof_pos = self.config.default_angles
for i in range(12):
motor_idx = self.config.joint2motor_idx[i]
self.low_cmd.motor_cmd[motor_idx].q = target_dof_pos[i]
self.low_cmd.motor_cmd[motor_idx].dq = 0.0
self.low_cmd.motor_cmd[motor_idx].kp = 20.0
self.low_cmd.motor_cmd[motor_idx].kd = 0.5
self.low_cmd.motor_cmd[motor_idx].tau = 0
self.send_cmd(self.low_cmd)
time.sleep(self.config.control_dt)
# === 调试：遥控器 & 模型输出 ===
# print(f"RC: lx={self.remote_controller.lx:+.2f} ly={self.remote_controller.ly:+.2f} "
# f"rx={self.remote_controller.rx:+.2f}")
# print(f"OBS cmd: {self.obs[6:9]}") # 遥控器信号在 obs 的位置
# print(f"RAW action: {self.action[:4]}...") # 只看前 4 个，防止刷屏
# print(f"TARGET Q: {target_dof_pos[::3]}") # 每 3 个关节抽 1 个，易读
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("net", type=str, help="network interface")
args = parser.parse_args()
config_path = f"{LEGGED_GYM_ROOT_DIR}/deploy/deploy_real/configs/go2.yaml"
config = Config(config_path)
ChannelFactoryInitialize(0, args.net)
controller = Controller(config)
controller.zero_torque_state()
controller.move_to_default_pos()
controller.default_pos_state()
while True:
try:
controller.run()
if controller.remote_controller.button[KeyMap.select] == 1:
break
except KeyboardInterrupt:
break
create_damping_cmd(controller.low_cmd)
controller.send_cmd(controller.low_cmd)
print('Exit')

BIN
deploy/pre_train/go2/go2_cts_150k.pt Normal file
View File

Binary file not shown.