diff --git a/.gitignore b/.gitignore index e144cbd..a2c090c 100644 --- a/.gitignore +++ b/.gitignore @@ -58,8 +58,11 @@ temp/ data/AMASS/Transitions/ data/g1_29dof/ data/smpl/ +data/ # Pickle files # *.pkl -homework.md \ No newline at end of file +homework.md + +assets/robot/gr3_6dof \ No newline at end of file diff --git a/README.md b/README.md index 72e1e02..0965e41 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,7 @@ Of course, you can download the other motion files you like from the [AMASS](htt ## Fitting motion data from AMASS First we need to fit the smpl shape to the humanoid robot: +export XDG_SESSION_TYPE=x11 ```bash uv run scripts/1-fit_smpl_shape.py ``` diff --git a/cfg/fourier_gr3_fitting.yaml b/cfg/fourier_gr3_fitting.yaml new file mode 100644 index 0000000..d11a97a --- /dev/null +++ b/cfg/fourier_gr3_fitting.yaml @@ -0,0 +1,58 @@ +humanoid_type: gr3_6dof + +asset: + assetFileName: "assets/robot/gr3_6dof/gr3_6dof/urdf/gr3_fourier_hand_6dof.urdf" + +extend_config: + - joint_name: "left_toe_link" + parent_name: "left_foot_roll_link" + pos: [0.1, 0.0, -0.02] + rot: [1.0, 0.0, 0.0, 0.0] + + - joint_name: "right_toe_link" + parent_name: "right_foot_roll_link" + pos: [0.1, 0.0, -0.02] + rot: [1.0, 0.0, 0.0, 0.0] + +base_link: "base_link" +joint_matches: + # - ["base_link", "Pelvis"] + - ["left_thigh_pitch_link", "L_Hip"] + - ["left_shank_pitch_link", "L_Knee"] + - ["left_foot_roll_link", "L_Ankle"] + - ["right_thigh_pitch_link", "R_Hip"] + - ["right_shank_pitch_link", "R_Knee"] + - ["right_foot_roll_link", "R_Ankle"] + - ["left_upper_arm_roll_link", "L_Shoulder"] + - ["left_lower_arm_pitch_link", "L_Elbow"] + - ["left_hand_roll_link", "L_Hand"] + - ["right_upper_arm_roll_link", "R_Shoulder"] + - ["right_lower_arm_pitch_link", "R_Elbow"] + - ["right_hand_roll_link", "R_Hand"] + - ["left_toe_link", "L_Toe"] + - ["right_toe_link", "R_Toe"] + + +smpl_pose_modifier: + Pelvis: "[np.pi/2, 0, np.pi/2]" + L_Shoulder: "[0, 0, -np.pi/2]" + R_Shoulder: "[0, 0, np.pi/2]" + # L_Elbow: "[0, -np.pi/2, 0]" + # R_Elbow: "[0, np.pi/2, 0]" + +quat_offset: + Pelvis: "[0.5, -0.5, -0.5, -0.5]" + L_Hip: "[0.5, -0.5, -0.5, -0.5]" + L_Knee: "[0.5, -0.5, -0.5, -0.5]" + L_Ankle: "[0.5, -0.5, -0.5, -0.5]" + R_Hip: "[0.5, -0.5, -0.5, -0.5]" + R_Knee: "[0.5, -0.5, -0.5, -0.5]" + R_Ankle: "[0.5, -0.5, -0.5, -0.5]" + L_Toe: "[0.5, -0.5, -0.5, -0.5]" + R_Toe: "[0.5, -0.5, -0.5, -0.5]" + L_Shoulder: "[0.7071, 0, -0.7071, 0]" + L_Elbow: "[-0.7071067811865475, 0.0, 0.7071067811865475, 0.0]" + L_Hand: "[-0.7071067811865475, 0.0, 0.7071067811865475, 0.0]" + R_Shoulder: "[0, 0.7071, 0, 0.7071]" + R_Elbow: "[0.0, 0.7071067811865475, 0.0, 0.7071067811865475]" + R_Hand: "[0.0, 0.7071067811865475, 0.0, 0.7071067811865475]" diff --git a/check_root.py b/check_root.py new file mode 100644 index 0000000..bd9e9dc --- /dev/null +++ b/check_root.py @@ -0,0 +1,18 @@ +import torch +import pytorch_kinematics as pk +import xml.etree.ElementTree as ET +import os + +path = "assets/robot/gr3_6dof/gr3_6dof/urdf/gr3_fourier_hand_6dof.urdf" +tree = ET.parse(path) +root = tree.getroot() +chain = pk.build_chain_from_urdf(ET.tostring(root, encoding='utf8')) + +print("Chain root:", chain.get_link_names()[0]) +print("Joint names:", chain.get_joint_parent_child_names()) + +th = torch.zeros([1, chain.n_joints]) +pos = chain.forward_kinematics(th) +print("Base link pose:", pos['base_link'].get_matrix()) +print("Right hip pose:", pos['right_thigh_pitch_link'].get_matrix()) +print("Left hip pose:", pos['left_thigh_pitch_link'].get_matrix()) diff --git a/check_smpl.py b/check_smpl.py new file mode 100644 index 0000000..4ce61da --- /dev/null +++ b/check_smpl.py @@ -0,0 +1,19 @@ +import torch +from smplx import SMPL +import os + +DATA_PATH = "data" +body_model = SMPL(model_path=os.path.join(DATA_PATH, "smpl"), gender="neutral") +betas = torch.zeros([1, 10]) +result = body_model(betas=betas) +joints = result.joints[0] # [24, 3] +pelvis = joints[0] +l_hip = joints[1] +r_hip = joints[2] +l_shoulder = joints[16] +r_shoulder = joints[17] + +print("SMPL Pelvis:", pelvis.tolist()) +print("SMPL L_Hip relative to Pelvis:", (l_hip - pelvis).tolist()) +print("SMPL R_Hip relative to Pelvis:", (r_hip - pelvis).tolist()) +print("SMPL L_Shoulder relative to Pelvis:", (l_shoulder - pelvis).tolist()) diff --git a/data/AMASS/SFU/0005/twistdance_kick_stageii.npz b/data/AMASS/SFU/0005/twistdance_kick_stageii.npz new file mode 100644 index 0000000..d1fff63 Binary files /dev/null and b/data/AMASS/SFU/0005/twistdance_kick_stageii.npz differ diff --git a/data/SFU.tar.bz2 b/data/SFU.tar.bz2 new file mode 100644 index 0000000..87fdc6f Binary files /dev/null and b/data/SFU.tar.bz2 differ diff --git a/meta.json b/meta.json index f352c74..4ed5775 100644 --- a/meta.json +++ b/meta.json @@ -1,72 +1,134 @@ { "body_names": [ - "pelvis", - "left_hip_pitch_link", - "left_hip_roll_link", - "left_hip_yaw_link", - "left_knee_link", - "left_ankle_pitch_link", - "left_ankle_roll_link", - "left_toe_link", - "right_hip_pitch_link", - "right_hip_roll_link", - "right_hip_yaw_link", - "right_knee_link", - "right_ankle_pitch_link", - "right_ankle_roll_link", - "right_toe_link", + "base_link", "waist_yaw_link", "waist_roll_link", + "waist_pitch_link", "torso_link", - "left_shoulder_pitch_link", - "left_shoulder_roll_link", - "left_shoulder_yaw_link", - "left_elbow_link", - "left_wrist_roll_link", - "left_wrist_pitch_link", - "left_wrist_yaw_link", - "left_rubber_hand", - "right_shoulder_pitch_link", - "right_shoulder_roll_link", - "right_shoulder_yaw_link", - "right_elbow_link", - "right_wrist_roll_link", - "right_wrist_pitch_link", - "right_wrist_yaw_link", - "right_rubber_hand", - "imu_in_torso", - "head_link", - "imu_in_pelvis" + "head_yaw_link", + "head_pitch_link", + "camera_link", + "right_upper_arm_pitch_link", + "right_upper_arm_roll_link", + "right_upper_arm_yaw_link", + "right_lower_arm_pitch_link", + "right_hand_yaw_link", + "right_hand_pitch_link", + "right_hand_roll_link", + "right_end_effector_link", + "R_hand_base_link", + "R_hand_wrist_link", + "R_thumb_proximal_yaw_link", + "R_thumb_proximal_pitch_link", + "R_thumb_distal_link", + "R_thumb_tip_link", + "R_index_proximal_link", + "R_index_intermediate_link", + "R_index_tip_link", + "R_middle_proximal_link", + "R_middle_intermediate_link", + "R_middle_tip_link", + "R_ring_proximal_link", + "R_ring_intermediate_link", + "R_ring_tip_link", + "R_pinky_proximal_link", + "R_pinky_intermediate_link", + "R_pinky_tip_link", + "left_upper_arm_pitch_link", + "left_upper_arm_roll_link", + "left_upper_arm_yaw_link", + "left_lower_arm_pitch_link", + "left_hand_yaw_link", + "left_hand_pitch_link", + "left_hand_roll_link", + "left_end_effector_link", + "L_hand_base_link", + "L_hand_wrist_link", + "L_thumb_proximal_yaw_link", + "L_thumb_proximal_pitch_link", + "L_thumb_distal_link", + "L_thumb_tip_link", + "L_index_proximal_link", + "L_index_intermediate_link", + "L_index_tip_link", + "L_middle_proximal_link", + "L_middle_intermediate_link", + "L_middle_tip_link", + "L_ring_proximal_link", + "L_ring_intermediate_link", + "L_ring_tip_link", + "L_pinky_proximal_link", + "L_pinky_intermediate_link", + "L_pinky_tip_link", + "right_thigh_pitch_link", + "right_thigh_roll_link", + "right_thigh_yaw_link", + "right_shank_pitch_link", + "right_foot_pitch_link", + "right_foot_roll_link", + "right_toe_link", + "left_thigh_pitch_link", + "left_thigh_roll_link", + "left_thigh_yaw_link", + "left_shank_pitch_link", + "left_foot_pitch_link", + "left_foot_roll_link", + "left_toe_link", + "imu_link" ], "joint_names": [ - "left_hip_pitch_joint", - "left_hip_roll_joint", - "left_hip_yaw_joint", - "left_knee_joint", - "left_ankle_pitch_joint", - "left_ankle_roll_joint", - "right_hip_pitch_joint", - "right_hip_roll_joint", - "right_hip_yaw_joint", - "right_knee_joint", - "right_ankle_pitch_joint", - "right_ankle_roll_joint", "waist_yaw_joint", "waist_roll_joint", "waist_pitch_joint", - "left_shoulder_pitch_joint", - "left_shoulder_roll_joint", - "left_shoulder_yaw_joint", - "left_elbow_joint", - "left_wrist_roll_joint", - "left_wrist_pitch_joint", - "left_wrist_yaw_joint", + "head_yaw_joint", + "head_pitch_joint", "right_shoulder_pitch_joint", "right_shoulder_roll_joint", "right_shoulder_yaw_joint", - "right_elbow_joint", - "right_wrist_roll_joint", + "right_elbow_pitch_joint", + "right_wrist_yaw_joint", "right_wrist_pitch_joint", - "right_wrist_yaw_joint" + "right_wrist_roll_joint", + "R_thumb_proximal_yaw_joint", + "R_thumb_proximal_pitch_joint", + "R_thumb_distal_joint", + "R_index_proximal_joint", + "R_index_intermediate_joint", + "R_middle_proximal_joint", + "R_middle_intermediate_joint", + "R_ring_proximal_joint", + "R_ring_intermediate_joint", + "R_pinky_proximal_joint", + "R_pinky_intermediate_joint", + "left_shoulder_pitch_joint", + "left_shoulder_roll_joint", + "left_shoulder_yaw_joint", + "left_elbow_pitch_joint", + "left_wrist_yaw_joint", + "left_wrist_pitch_joint", + "left_wrist_roll_joint", + "L_thumb_proximal_yaw_joint", + "L_thumb_proximal_pitch_joint", + "L_thumb_distal_joint", + "L_index_proximal_joint", + "L_index_intermediate_joint", + "L_middle_proximal_joint", + "L_middle_intermediate_joint", + "L_ring_proximal_joint", + "L_ring_intermediate_joint", + "L_pinky_proximal_joint", + "L_pinky_intermediate_joint", + "right_hip_pitch_joint", + "right_hip_roll_joint", + "right_hip_yaw_joint", + "right_knee_pitch_joint", + "right_ankle_pitch_joint", + "right_ankle_roll_joint", + "left_hip_pitch_joint", + "left_hip_roll_joint", + "left_hip_yaw_joint", + "left_knee_pitch_joint", + "left_ankle_pitch_joint", + "left_ankle_roll_joint" ] } \ No newline at end of file diff --git a/scripts/1-fit_smpl_shape.py b/scripts/1-fit_smpl_shape.py index d7d0baf..21d3e8b 100644 --- a/scripts/1-fit_smpl_shape.py +++ b/scripts/1-fit_smpl_shape.py @@ -40,40 +40,65 @@ def build_chain(cfg) -> pk.Chain: - mjcf_path = cfg.asset.assetFileName + path = cfg.asset.assetFileName + is_urdf = path.endswith(".urdf") - tree = ET.parse(mjcf_path) + tree = ET.parse(path) root = tree.getroot() - # remove the free joint of the base link - root_name = cfg.get("root_name", "pelvis") - root_body = root.find(f".//body[@name='{root_name}']") - root_joint = root.find(".//joint[@type='free']") - root_body.remove(root_joint) + if not is_urdf: # MJCF specific: remove the free joint of the base link + root_name = cfg.get("root_name", "pelvis") + root_body = root.find(f".//body[@name='{root_name}']") + if root_body is not None: + root_joint = root_body.find(".//joint[@type='free']") + if root_joint is not None: + root_body.remove(root_joint) for extend_config in cfg.extend_config: - parent = root.find(f".//body[@name='{extend_config.parent_name}']") + if is_urdf: + parent = root.find(f".//link[@name='{extend_config.parent_name}']") + else: + parent = root.find(f".//body[@name='{extend_config.parent_name}']") + if parent is None: - raise ValueError(f"Parent body {extend_config.parent_name} not found in MJCF") + raise ValueError(f"Parent {extend_config.parent_name} not found") pos = extend_config.pos - rot = extend_config.rot - # create and insert a dummy body with a fixed joint - body = ET.Element("body", name=extend_config.joint_name) - body.set("pos", f"{pos[0]} {pos[1]} {pos[2]}") - body.set("quat", f"{rot[0]} {rot[1]} {rot[2]} {rot[3]}") - inertial = ET.Element("inertial", pos="0 0 0", quat="0 0 0 1", mass="0.1", diaginertia="0.1 0.1 0.1") - body.append(inertial) - parent.append(body) + rot = extend_config.rot # [w, x, y, z] + + if is_urdf: # URDF: Insert link and fixed joint + link_name = extend_config.joint_name + ET.SubElement(root, "link", name=link_name) + + joint = ET.SubElement(root, "joint", name=f"fixed_{link_name}", type="fixed") + ET.SubElement(joint, "parent", link=extend_config.parent_name) + ET.SubElement(joint, "child", link=link_name) + + # URDF origin uses rpy, rot in yaml is quaternion [w, x, y, z] + # scipy needs [x, y, z, w] + r = sRot.from_quat([rot[1], rot[2], rot[3], rot[0]]) + rpy = r.as_euler('xyz') + ET.SubElement(joint, "origin", xyz=f"{pos[0]} {pos[1]} {pos[2]}", rpy=f"{rpy[0]} {rpy[1]} {rpy[2]}") + else: # MJCF: Insert body + body = ET.Element("body", name=extend_config.joint_name) + body.set("pos", f"{pos[0]} {pos[1]} {pos[2]}") + body.set("quat", f"{rot[0]} {rot[1]} {rot[2]} {rot[3]}") + inertial = ET.Element("inertial", pos="0 0 0", quat="0 0 0 1", mass="0.1", diaginertia="0.1 0.1 0.1") + body.append(inertial) + parent.append(body) cwd = os.getcwd() - os.chdir(os.path.dirname(mjcf_path)) - chain = pk.build_chain_from_mjcf(ET.tostring(root, method="xml"), body=root_name) + os.chdir(os.path.dirname(path)) + if is_urdf: + chain = pk.build_chain_from_urdf(ET.tostring(root, encoding='utf8')) + else: + chain = pk.build_chain_from_mjcf(ET.tostring(root, method="xml"), body=root_name) os.chdir(cwd) return chain -@hydra.main(version_base=None, config_path="../cfg", config_name="unitree_g1_fitting") +@hydra.main(version_base=None, config_path="../cfg", config_name="fourier_gr3_fitting") +# @hydra.main(version_base=None, config_path="../cfg", config_name="unitree_g1_fitting") def main(cfg): chain = build_chain(cfg) chain.print_tree() @@ -139,6 +164,12 @@ def init_mesh(vertices, color=[0.3, 0.3, 0.3]): pelvis = result.joints[:, None, 0] smpl_keypoints = result.joints[:, smpl_joint_idx] - pelvis + #加入正则化 + # mse_loss = (robot_keypoints - smpl_keypoints).square().sum() + # reg_loss = (betas ** 2).mean() + # loss = mse_loss + 0.01 * reg_loss + + loss = (robot_keypoints - smpl_keypoints).square().sum() opt.zero_grad() loss.backward() @@ -148,7 +179,7 @@ def init_mesh(vertices, color=[0.3, 0.3, 0.3]): v = result.vertices[0] - result.joints[:, 0] vertices.append(v) smpl_keypoints_history.append(smpl_keypoints[0].detach()) - print(f"iter {i}, loss: {loss.item():.3f}") + print(f"iter {i}, loss: {loss.item():.3f}, betas: {betas.data}") vis = o3d.visualization.Visualizer() vis.create_window() @@ -164,7 +195,7 @@ def init_mesh(vertices, color=[0.3, 0.3, 0.3]): vis.add_geometry(mesh) pcd = o3d.geometry.PointCloud() - pcd.points = o3d.utility.Vector3dVector(smpl_keypoints_history[i] + offset + torch.tensor([1., 0., 0.])) + pcd.points = o3d.utility.Vector3dVector(smpl_keypoints_history[i] + offset) pcd.colors = o3d.utility.Vector3dVector([[0, 0, 1] for _ in range(len(smpl_keypoints_history[i]))]) vis.add_geometry(pcd) @@ -187,8 +218,9 @@ def init_mesh(vertices, color=[0.3, 0.3, 0.3]): # increase point size for better visualization opt = vis.get_render_option() - point_size = 10.0 # You can adjust this value according to your needs - opt.point_size = point_size + if opt: + point_size = 10.0 + opt.point_size = point_size vis.run() diff --git a/scripts/3-fit_smpl_motion_mp.py b/scripts/3-fit_smpl_motion_mp.py index 5d92d96..bbdbc71 100644 --- a/scripts/3-fit_smpl_motion_mp.py +++ b/scripts/3-fit_smpl_motion_mp.py @@ -60,42 +60,88 @@ def parse_joint_limits_from_mjcf(root): return limits +def parse_joint_limits_from_urdf(root): + limits = {} + for j in root.findall(".//joint"): + jtype = j.get("type") + if jtype not in ("revolute", "prismatic"): + continue + name = j.get("name") + limit = j.find("limit") + if name is None or limit is None: + continue + low = float(limit.get("lower", -3.14159265359)) + high = float(limit.get("upper", 3.14159265359)) + limits[name] = (low, high) + return limits + + def build_chain(cfg) -> pk.Chain: - mjcf_path = cfg.asset.assetFileName + path = cfg.asset.assetFileName + is_urdf = path.endswith(".urdf") - tree = ET.parse(mjcf_path) + tree = ET.parse(path) root = tree.getroot() - # remove the free joint of the base link - root_name = cfg.get("root_name", "pelvis") - root_body = root.find(f".//body[@name='{root_name}']") - root_joint = root.find(".//joint[@type='free']") - root_body.remove(root_joint) - root_body.set("pos", "0 0 0") + if is_urdf: + joint_limits = parse_joint_limits_from_urdf(root) + else: + # MJCF specific: remove the free joint of the base link + root_name = cfg.get("root_name", "pelvis") + root_body = root.find(f".//body[@name='{root_name}']") + if root_body is not None: + root_joint = root_body.find(".//joint[@type='free']") + if root_joint is not None: + root_body.remove(root_joint) + root_body.set("pos", "0 0 0") + joint_limits = parse_joint_limits_from_mjcf(root) for extend_config in cfg.extend_config: - parent = root.find(f".//body[@name='{extend_config.parent_name}']") + if is_urdf: + parent = root.find(f".//link[@name='{extend_config.parent_name}']") + else: + parent = root.find(f".//body[@name='{extend_config.parent_name}']") + if parent is None: - raise ValueError(f"Parent body {extend_config.parent_name} not found in MJCF") + raise ValueError(f"Parent {extend_config.parent_name} not found") pos = extend_config.pos - rot = extend_config.rot - # create and insert a dummy body with a fixed joint - body = ET.Element("body", name=extend_config.joint_name) - body.set("pos", f"{pos[0]} {pos[1]} {pos[2]}") - body.set("quat", f"{rot[0]} {rot[1]} {rot[2]} {rot[3]}") - inertial = ET.Element("inertial", pos="0 0 0", quat="0 0 0 1", mass="0.1", diaginertia="0.1 0.1 0.1") - body.append(inertial) - parent.append(body) + rot = extend_config.rot # [w, x, y, z] + + if is_urdf: # URDF: Insert link and fixed joint + link_name = extend_config.joint_name + ET.SubElement(root, "link", name=link_name) + + joint = ET.SubElement(root, "joint", name=f"fixed_{link_name}", type="fixed") + ET.SubElement(joint, "parent", link=extend_config.parent_name) + ET.SubElement(joint, "child", link=link_name) + + # URDF origin uses rpy, rot in yaml is quaternion [w, x, y, z] + # scipy needs [x, y, z, w] + r = sRot.from_quat([rot[1], rot[2], rot[3], rot[0]]) + rpy = r.as_euler('xyz') + ET.SubElement(joint, "origin", xyz=f"{pos[0]} {pos[1]} {pos[2]}", rpy=f"{rpy[0]} {rpy[1]} {rpy[2]}") + else: # MJCF: Insert body + body = ET.Element("body", name=extend_config.joint_name) + body.set("pos", f"{pos[0]} {pos[1]} {pos[2]}") + body.set("quat", f"{rot[0]} {rot[1]} {rot[2]} {rot[3]}") + inertial = ET.Element("inertial", pos="0 0 0", quat="0 0 0 1", mass="0.1", diaginertia="0.1 0.1 0.1") + body.append(inertial) + parent.append(body) cwd = os.getcwd() - os.chdir(os.path.dirname(mjcf_path)) - chain = pk.build_chain_from_mjcf(ET.tostring(root, method="xml"), body=root_name) + os.chdir(os.path.dirname(path)) + if is_urdf: + chain = pk.build_chain_from_urdf(ET.tostring(root, encoding='utf8')) + else: + chain = pk.build_chain_from_mjcf(ET.tostring(root, method="xml"), body=root_name) os.chdir(cwd) - - joint_limits = parse_joint_limits_from_mjcf(root) - return chain, joint_limits - + + # Filter joint_limits to only include joints present in the chain + parameter_names = chain.get_joint_parameter_names() + filtered_limits = [joint_limits[name] for name in parameter_names] + + return chain, filtered_limits def lerp(x, xp, fp): return np.stack([np.interp(x, xp, fp[:, i]) for i in range(fp.shape[1])], axis=1) @@ -139,10 +185,9 @@ def fit_motion(cfg, motion_path: str, fitted_shape: torch.Tensor): assert len(joint_names) == len(joint_limits), f"Number of joints in chain ({len(joint_names)}) does not match number of joint limits ({len(joint_limits)})" low_list, high_list = [], [] - for joint_name, (joint_limit_key, joint_limit_val) in zip(joint_names, joint_limits.items()): - assert joint_name == joint_limit_key, f"Joint name {joint_name} does not match joint name in joint limits ({joint_limit_key})" - low_list.append(joint_limit_val[0]) - high_list.append(joint_limit_val[1]) + for joint_name, (l, h) in zip(joint_names, joint_limits): + low_list.append(l) + high_list.append(h) low = torch.as_tensor(low_list, dtype=torch.float32).unsqueeze(0) # [1, J] high = torch.as_tensor(high_list, dtype=torch.float32).unsqueeze(0) # [1, J] @@ -246,8 +291,8 @@ def mat_rotate(rotmat, v): opt.zero_grad() loss.backward() opt.step() - # if i % 10 == 0: - # print(f"iter {i}, loss {100 * loss.item():.3f}") + if i % 490 == 0: + print(f"iter {i}, loss {100 * loss.item():.3f}") with torch.no_grad(): robot_keypoints_b = torch.stack([ @@ -274,7 +319,9 @@ def _worker(args): cfg, path, betas = args return fit_motion(cfg, path, betas) -@hydra.main(version_base=None, config_path="../cfg", config_name="unitree_g1_fitting") +# @hydra.main(version_base=None, config_path="../cfg", config_name="unitree_g1_fitting") +@hydra.main(version_base=None, config_path="../cfg", config_name="fourier_gr3_fitting") + def main(cfg): if os.path.isdir(cfg.motion_path): motion_paths = glob.glob(os.path.join(cfg.motion_path, "**/*.npz"), recursive=True) diff --git a/scripts/4-vis_mj.py b/scripts/4-vis_mj.py index 30767a8..7a6e1b2 100644 --- a/scripts/4-vis_mj.py +++ b/scripts/4-vis_mj.py @@ -1,68 +1,130 @@ import os import numpy as np import mujoco -import mujoco.viewer import hydra +from hydra.utils import get_original_cwd import time -import glob import threading -from scipy.spatial.transform import Rotation as sRot, Slerp import joblib +import viser +import viser.transforms as tf +from viser.extras import ViserUrdf +import yourdfpy DATA_PATH = os.path.join(os.path.dirname(__file__), "../data") - class MotionVis: - def __init__(self, motions): + def __init__(self, motions, server, viser_urdf, robot_frame, index_map): self.motions = motions self.i = 0 self.t = 0 self.motion = self.motions[self.i] + self.server = server + self.viser_urdf = viser_urdf + self.robot_frame = robot_frame + self.index_map = index_map - def key_callback(self, keycode): - # print(keycode) - if keycode == 265: # up arrow, prev motion - self.i = (self.i - 1) % len(self.motions) - self.motion = self.motions[self.i] - self.t = 0 - elif keycode == 264: # down arrow, next motion - self.i = (self.i + 1) % len(self.motions) - self.motion = self.motions[self.i] - self.t = 0 - - def run(self, model, data: mujoco.MjData, viewer: mujoco.viewer.Handle): - while viewer.is_running(): - with viewer.lock(): - data.qpos[7:] = self.motion["joint_pos"][self.t] - data.qpos[:3] = self.motion["root_pos_w"][self.t] - data.qpos[3:7] = self.motion["root_quat_w"][self.t][[3, 0, 1, 2]] - self.t = (self.t + 1) % self.motion["joint_pos"].shape[0] - mujoco.mj_forward(model, data) - time.sleep(1 / self.motion["fps"]) + # adding GUI controls + with self.server.gui.add_folder("Motion Control"): + self.gui_motion_name = self.server.gui.add_text("Current", initial_value="Motion 0") + prev_btn = self.server.gui.add_button("Previous (Up)") + next_btn = self.server.gui.add_button("Next (Down)") + + @prev_btn.on_click + def _(_) : self.change_motion(-1) + @next_btn.on_click + def _(_) : self.change_motion(1) + + def change_motion(self, delta): + self.i = (self.i + delta) % len(self.motions) + self.motion = self.motions[self.i] + self.t = 0 + self.gui_motion_name.value = f"Motion {self.i}" + def run(self, model, data: mujoco.MjData): + n_joint = self.motion["joint_pos"].shape[1] + + has_free_joint = (model.nq - n_joint) == 7 # 7 DOF for free joint + while True: + + joint_angles = self.motion["joint_pos"][self.t] + root_pos = self.motion["root_pos_w"][self.t] + q = self.motion["root_quat_w"][self.t] # [x, y, z, w] + + if has_free_joint: + data.qpos[7:] = joint_angles + data.qpos[:3] = root_pos + data.qpos[3:7] = q[[3, 0, 1, 2]] + + mujoco.mj_forward(model, data) + + # update robot frame + self.robot_frame.position = root_pos + self.robot_frame.wxyz = (q[3], q[0], q[1], q[2]) + + self.viser_urdf.update_cfg(joint_angles[self.index_map]) + + self.t = (self.t + 1) % self.motion["joint_pos"].shape[0] + time.sleep(1 / self.motion["fps"]) -@hydra.main(version_base=None, config_path="../cfg", config_name="unitree_g1_fitting") +@hydra.main(version_base=None, config_path="../cfg", config_name="fourier_gr3_fitting") +# @hydra.main(version_base=None, config_path="../cfg", config_name="unitree_g1_fitting") def main(cfg): - mjcf_path = cfg.asset.assetFileName - model = mujoco.MjModel.from_xml_path(mjcf_path) + asset_path = cfg.asset.assetFileName + if not os.path.isabs(asset_path): + asset_path = os.path.join(get_original_cwd(), asset_path) + + print(f"Loading URDF from: {asset_path}") + urdf_model = yourdfpy.URDF.load(asset_path) + + + model = mujoco.MjModel.from_xml_path(asset_path) data = mujoco.MjData(model) + # launch viser server + server = viser.ViserServer() + server.scene.add_grid("/ground", width=5.0, height=5.0) + + robot_frame = server.scene.add_frame("/robot") + viser_urdf = ViserUrdf(server, urdf_model, root_node_name="/robot") + + # build pk chain for index mapping + import pytorch_kinematics as pk + with open(asset_path, "rb") as f: + chain = pk.build_chain_from_urdf(f.read()) + pk_joint_names = chain.get_joint_parameter_names() + urdf_joint_names = list(viser_urdf.get_actuated_joint_limits().keys()) + index_map = [pk_joint_names.index(name) for name in urdf_joint_names] + + # load motions and start visualization motions = joblib.load(cfg.motion_file) motions = list(motions.values()) - motion_vis = MotionVis(motions) - - viewer = mujoco.viewer.launch_passive(model, data, key_callback=motion_vis.key_callback) + # Check for joints that are always zero + all_joint_pos = np.concatenate([m["joint_pos"] for m in motions], axis=0) + is_always_zero = np.all(np.abs(all_joint_pos) < 1e-6, axis=0) + zero_indices = np.where(is_always_zero)[0] + nonzero_indices = np.where(~is_always_zero)[0] + - def render_async(): - while viewer.is_running(): - viewer.sync() - time.sleep(0.02) - threading.Thread(target=render_async).start() - motion_vis.run(model, data, viewer) + print("\n" + "="*50) + print(f"检测到 {len(zero_indices)}/{len(pk_joint_names)} 个关节在所有序列中均为 0 (无运动):") + for idx in zero_indices: + print(f"编号: {idx:2d} | 关节名: {pk_joint_names[idx]}") + print("="*50 + "\n") + print("-" * 20) + print(f"检测到 {len(nonzero_indices)}/{len(pk_joint_names)} 个关节具有有效运动:") + for idx in nonzero_indices: + print(f"编号: {idx:2d} | 关节名: {pk_joint_names[idx]}") + print("="*50 + "\n") + + motion_vis = MotionVis(motions, server, viser_urdf, robot_frame, index_map) + + print(f"Viser server started. Please open your browser to view.") + motion_vis.run(model, data) if __name__ == "__main__": main() \ No newline at end of file diff --git a/scripts/4-vis_xml.py b/scripts/4-vis_xml.py new file mode 100644 index 0000000..30767a8 --- /dev/null +++ b/scripts/4-vis_xml.py @@ -0,0 +1,68 @@ +import os +import numpy as np +import mujoco +import mujoco.viewer +import hydra +import time +import glob +import threading +from scipy.spatial.transform import Rotation as sRot, Slerp +import joblib + +DATA_PATH = os.path.join(os.path.dirname(__file__), "../data") + + +class MotionVis: + def __init__(self, motions): + self.motions = motions + self.i = 0 + self.t = 0 + self.motion = self.motions[self.i] + + def key_callback(self, keycode): + # print(keycode) + if keycode == 265: # up arrow, prev motion + self.i = (self.i - 1) % len(self.motions) + self.motion = self.motions[self.i] + self.t = 0 + elif keycode == 264: # down arrow, next motion + self.i = (self.i + 1) % len(self.motions) + self.motion = self.motions[self.i] + self.t = 0 + + def run(self, model, data: mujoco.MjData, viewer: mujoco.viewer.Handle): + while viewer.is_running(): + with viewer.lock(): + data.qpos[7:] = self.motion["joint_pos"][self.t] + data.qpos[:3] = self.motion["root_pos_w"][self.t] + data.qpos[3:7] = self.motion["root_quat_w"][self.t][[3, 0, 1, 2]] + self.t = (self.t + 1) % self.motion["joint_pos"].shape[0] + mujoco.mj_forward(model, data) + time.sleep(1 / self.motion["fps"]) + + +@hydra.main(version_base=None, config_path="../cfg", config_name="unitree_g1_fitting") +def main(cfg): + mjcf_path = cfg.asset.assetFileName + + model = mujoco.MjModel.from_xml_path(mjcf_path) + data = mujoco.MjData(model) + + motions = joblib.load(cfg.motion_file) + motions = list(motions.values()) + + motion_vis = MotionVis(motions) + + viewer = mujoco.viewer.launch_passive(model, data, key_callback=motion_vis.key_callback) + + def render_async(): + while viewer.is_running(): + viewer.sync() + time.sleep(0.02) + threading.Thread(target=render_async).start() + + motion_vis.run(model, data, viewer) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/scripts/g1_29dof_shape.npz b/scripts/g1_29dof_shape.npz index 67f4f85..4a0a061 100644 Binary files a/scripts/g1_29dof_shape.npz and b/scripts/g1_29dof_shape.npz differ diff --git a/scripts/gr3_6dof_shape.npz b/scripts/gr3_6dof_shape.npz new file mode 100644 index 0000000..d2f1af3 Binary files /dev/null and b/scripts/gr3_6dof_shape.npz differ diff --git a/scripts/inspect_robot.py b/scripts/inspect_robot.py new file mode 100644 index 0000000..bd057ba --- /dev/null +++ b/scripts/inspect_robot.py @@ -0,0 +1,55 @@ +import os +import torch +import yaml +import numpy as np +import pytorch_kinematics as pk +import xml.etree.ElementTree as ET +from scipy.spatial.transform import Rotation as sRot + +def build_chain(cfg) -> pk.Chain: + path = cfg['asset']['assetFileName'] + is_urdf = path.endswith(".urdf") + tree = ET.parse(path) + root = tree.getroot() + cwd = os.getcwd() + os.chdir(os.path.dirname(path)) + chain = pk.build_chain_from_urdf(ET.tostring(root, encoding='utf8')) + os.chdir(cwd) + return chain + +with open("cfg/fourier_gr3_fitting.yaml", "r") as f: + cfg = yaml.safe_load(f) + +chain = build_chain(cfg) +th = torch.zeros([1, chain.n_joints]) +robot_body_pose = chain.forward_kinematics(th) + +robot_body_names = [m[0] for m in cfg['joint_matches']] +robot_keypoints = [] +for name in robot_body_names: + if name in robot_body_pose: + robot_keypoints.append(robot_body_pose[name].get_matrix()[:, :3, 3]) + else: + print(f"Warning: {name} not found in robot_body_pose") + robot_keypoints.append(torch.zeros([1, 3])) +robot_keypoints = torch.stack(robot_keypoints, dim=1) +root_translation = robot_body_pose[chain.get_link_names()[0]].get_matrix()[:, :3, 3] +robot_keypoints = robot_keypoints - root_translation.unsqueeze(1) + +print("Robot Keypoints (centered at root):") +for i, name in enumerate(robot_body_names): + print(f"{name}: {robot_keypoints[0, i].tolist()}") + +# Calculate distances +l_hip = robot_keypoints[0, robot_body_names.index("left_thigh_pitch_link")] +r_hip = robot_keypoints[0, robot_body_names.index("right_thigh_pitch_link")] +hip_width = torch.norm(l_hip - r_hip).item() +print(f"Hip width: {hip_width:.4f}m") + +l_knee = robot_keypoints[0, robot_body_names.index("left_shank_pitch_link")] +thigh_length = torch.norm(l_hip - l_knee).item() +print(f"Thigh length: {thigh_length:.4f}m") + +l_ankle = robot_keypoints[0, robot_body_names.index("left_foot_roll_link")] +shank_length = torch.norm(l_knee - l_ankle).item() +print(f"Shank length: {shank_length:.4f}m")