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

v1.1.1-rc2; add README

Browse Source
This commit is contained in:
wty-yy
2026-01-27 21:36:44 +08:00
parent 064d0c90ea
commit 93c815b2d9
10 changed files with 468 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
deploy/deploy_mujoco/deploy_go2.py
View File

@@ -1,3 +1,10 @@
import sys
from pathlib import Path
PATH_PARENT = Path(__file__).parent
sys.path.append(str(PATH_PARENT))
from utils import MujocoRenderUtils
import os
import time
import mujoco.viewer
import mujoco
@@ -7,12 +14,10 @@ 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]
@@ -110,7 +115,7 @@ if __name__ == "__main__":
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")
video_save_dir = str(PATH_PARENT / "videos")
os.makedirs(video_save_dir, exist_ok=True)
model_name = os.path.basename(policy_path).split('.')[0]
@@ -134,17 +139,18 @@ if __name__ == "__main__":
# load policy
policy = torch.jit.load(policy_path)
video_fps = 50
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}")
mujoco_render_utils = MujocoRenderUtils(video_fps, m.opt.timestep)
if visualize_moe_weights:
plt.ion()
@@ -153,7 +159,7 @@ if __name__ == "__main__":
bars = None
if save_moe_latent:
latent_save_dir = str(Path(__file__).parent / "data_latents")
latent_save_dir = str(PATH_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)
@@ -164,9 +170,9 @@ if __name__ == "__main__":
# 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
viewer.cam.distance = 2.0
viewer.cam.elevation = -20.0
viewer.cam.azimuth = 60.0
# Close the viewer automatically after simulation_duration wall-seconds.
start = time.time()
@@ -188,10 +194,12 @@ if __name__ == "__main__":
# 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)
mujoco_render_utils.update(cmd, d)
if save_video and counter % frame_skip == 0:
try:
renderer.update_scene(d, camera=viewer.cam)
mujoco_render_utils.update_external_rendering(renderer, ctype='renderer')
frame = renderer.render()
writer.append_data(frame)
except Exception as e:
@@ -249,6 +257,7 @@ if __name__ == "__main__":
target_dof_pos = action * action_scale + default_angles
# Pick up changes to the physics state, apply perturbations, update options from GUI.
mujoco_render_utils.update_external_rendering(viewer, ctype='viewer')
viewer.sync()
# Rudimentary time keeping, will drift relative to wall clock.

112
deploy/deploy_mujoco/utils.py Normal file
View File

@@ -0,0 +1,112 @@
from typing import Union, Literal
import numpy as np
import mujoco
import mujoco.viewer
class MujocoRenderUtils:
def __init__(self, render_fps, sim_dt):
self.target_velocity = None
self.vis_smooth_factor = 1.0
self.ren_smooth_factor = 1.0
self.vis_cur_vel = np.zeros(3)
self.ren_cur_vel = np.zeros(3)
self.mj_data = None
def update(self, target_velocity, mj_data):
self.target_velocity = target_velocity
self.mj_data = mj_data
def update_external_rendering(self,
handle: Union[mujoco.viewer.Handle, mujoco.Renderer],
ctype: Literal['viewer', 'renderer'],
):
""" Update external rendering handle (viewer or renderer). """
def add_thick_arrow(geom_elem, pos, vec, rgba, scale=0.7):
vel_norm = np.linalg.norm(vec)
display_norm = min(vel_norm * scale, 1.0)
if display_norm < 0.10:
mujoco.mjv_initGeom(
geom_elem,
type=mujoco.mjtGeom.mjGEOM_NONE,
size=[0,0,0], pos=pos, mat=np.eye(3).flatten(), rgba=[0,0,0,0]
)
return
mat = np.zeros(9)
target_quat = np.zeros(4)
vec_normalized = vec / vel_norm
mujoco.mju_quatZ2Vec(target_quat, vec_normalized)
mujoco.mju_quat2Mat(mat, target_quat)
mat = mat.reshape(3, 3)
mat[:, 2] *= display_norm
mujoco.mjv_initGeom(
geom_elem,
type=mujoco.mjtGeom.mjGEOM_ARROW,
size=[0.02, 0.02, display_norm], # [height, width, length]
pos=pos,
mat=mat.flatten(),
rgba=rgba
)
viewer_geom_idx = 0
if ctype == 'viewer':
handle.user_scn.ngeom = 0 # reset user scene geometry
if self.target_velocity is not None:
base_pos_world = self.mj_data.qpos[:3]
base_quat = self.mj_data.qpos[3:7]
# rendering arrows start position
offset_body = np.array([0.0, 0.0, 0.2])
offset_world = np.zeros(3)
mujoco.mju_rotVecQuat(offset_world, offset_body, base_quat)
start_pos = base_pos_world + offset_world
tgt_vel_body = np.array([self.target_velocity[0], self.target_velocity[1], 0.0])
raw_cur_vel_world = self.mj_data.qvel[:3]
raw_cur_vel = np.zeros(3)
neg_quat = np.zeros(4)
mujoco.mju_negQuat(neg_quat, base_quat)
mujoco.mju_rotVecQuat(raw_cur_vel, raw_cur_vel_world, neg_quat)
cur_vel_body = np.array([raw_cur_vel[0], raw_cur_vel[1], 0.0])
# EMA: v_smooth = alpha * v_new + (1 - alpha) * v_old
# alpha = self.vis_smooth_factor if ctype == 'viewer' else self.ren_smooth_factor
self.vis_cur_vel = cur_vel_body
self.ren_cur_vel = cur_vel_body
tgt_vel_world = np.zeros(3)
cur_vel_world = np.zeros(3)
mujoco.mju_rotVecQuat(tgt_vel_world, tgt_vel_body, base_quat)
if ctype == 'viewer':
mujoco.mju_rotVecQuat(cur_vel_world, self.vis_cur_vel, base_quat)
else:
mujoco.mju_rotVecQuat(cur_vel_world, self.ren_cur_vel, base_quat)
COLOR_CMD = [0, 1, 0, 1] # Green 0x00ff00
COLOR_REAL = [0, 0, 1, 1] # Blue 0x0000ff
if ctype == 'viewer':
# Cmd Arrow
add_thick_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, tgt_vel_world, COLOR_CMD)
viewer_geom_idx += 1
# Real Arrow
add_thick_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, cur_vel_world, COLOR_REAL)
viewer_geom_idx += 1
else:
# Renderer Append
handle.scene.ngeom += 1
add_thick_arrow(handle.scene.geoms[handle.scene.ngeom - 1], start_pos, tgt_vel_world, COLOR_CMD)
handle.scene.ngeom += 1
add_thick_arrow(handle.scene.geoms[handle.scene.ngeom - 1], start_pos, cur_vel_world, COLOR_REAL)
if ctype == 'viewer':
handle.user_scn.ngeom = viewer_geom_idx

3
deploy/deploy_real/config_go2.py
View File

@@ -1,4 +1,5 @@
from legged_gym import LEGGED_GYM_ROOT_DIR
from pathlib import Path
LEGGED_GYM_ROOT_DIR = str(Path(__file__).parents[2])
import numpy as np
import yaml

2
deploy/deploy_real/configs/go2.yaml
View File

@@ -6,7 +6,7 @@ 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"
policy_path: "{LEGGED_GYM_ROOT_DIR}/deploy/pre_train/go2/go2_moe_cts_137000_0.6365.pt"
joint2motor_idx: [3,4,5,0,1,2,9,10,11,6,7,8]

16
deploy/deploy_real/deploy_real_go2.py
View File

@@ -1,4 +1,5 @@
from legged_gym import LEGGED_GYM_ROOT_DIR
from pathlib import Path
LEGGED_GYM_ROOT_DIR = str(Path(__file__).parents[2])
import numpy as np
import time
import torch
@@ -85,11 +86,13 @@ class Controller:
def zero_torque_state(self):
print("Enter zero torque state.")
print("Waiting for the start signal...")
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)
print("Start signal received.")
print("Press *select* button to exit.")
def move_to_default_pos(self):
@@ -121,7 +124,7 @@ class Controller:
def default_pos_state(self):
print("Enter default pos state.")
print("Waiting for the Button A signal...")
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]
@@ -162,7 +165,12 @@ class Controller:
self.obs[33:45] = self.action
obs_tensor = torch.from_numpy(self.obs).unsqueeze(0)
self.action = self.policy(obs_tensor).detach().numpy().squeeze()
results = self.policy(obs_tensor)
if isinstance(results, tuple):
self.action = results[0]
else:
self.action = results
self.action = self.action.detach().numpy().squeeze()
target_dof_pos = self.config.default_angles + self.action * self.config.action_scale
# target_dof_pos = self.config.default_angles