NMR / tools /data_process /process /process_seed_g1.py

upload

45950ff 5 months ago

9.51 kB

import sys
sys.path.append('/mnt/shenzhen2cephfs/capybarali/codes/humanoid')

import torch, yaml, os
from tqdm import tqdm
from src.utils.rotation_conversions import quaternion_to_matrix, matrix_to_rotation_6d, matrix_to_axis_angle, rotation_6d_to_matrix, matrix_to_quaternion, axis_angle_to_matrix
import numpy as np
from argparse import ArgumentParser
import joblib
from data.vis import vis_3d_motion
from data.vis_g1 import vis_3d_g1
from copy import deepcopy
from scipy.spatial.transform import Rotation as R
from scipy.spatial.transform import Slerp
from scipy import interpolate
from joblib import Parallel, delayed
import argparse
import numpy as np
import pandas as pd
from scipy.spatial.transform import Rotation as R
import pinocchio as pin

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

parser = ArgumentParser(description="Launch MoCap processing")
parser.add_argument('--save_root', type=str, default="data/seed_g1_217")
parser.add_argument('--start_idx', type=int, default=0)
parser.add_argument('--interval', type=int, default=1)
args = parser.parse_args()

os.makedirs(args.save_root, exist_ok=True)   

def extract_g1_component(x):
    vel_xy = x[:, :2]
    dof = x[:, -29:]
    root_rot_mat = rotation_6d_to_matrix(x[:, 2:8])
    trans_xyz = torch.cat([
        torch.cumsum(vel_xy[:, :1], dim=0),  # new_x = cumsum(Δnew_x)
        x[:, 9:10],                           # new_y = height (direct)
        torch.cumsum(vel_xy[:, 1:], dim=0),  # new_z = cumsum(Δnew_z)
    ], dim=1)

    rot_mat = torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]]).float()
    global_orient_mat = root_rot_mat.squeeze(1).float()
    global_orient_mat = torch.einsum('ij,tjk->tik', rot_mat, global_orient_mat)
    rot_quat = matrix_to_quaternion(global_orient_mat)  # (T, 4) wxyz order

    transl = trans_xyz.float()
    transl = torch.einsum('ij,tj->ti', rot_mat, transl)

    return dof, rot_quat, transl

def get_g1_motion(joint_pos, body_pos_w, body_quat_w):
    dof = joint_pos # T, 29
    root_ori = body_quat_w[:, 0] # T, 4, wxyz
    joints = body_pos_w # T, J, 3
    # import ipdb; ipdb.set_trace()
    # 插值
    rotation_matrix = torch.tensor([[1.0, 0, 0], [0, 0, -1], [0, 1, 0]]).inverse()

    global_orient_mat = quaternion_to_matrix(torch.from_numpy(root_ori)).float()
    global_orient_mat = torch.einsum('ij,tjk->tik', rotation_matrix, global_orient_mat)
    global_orient = matrix_to_axis_angle(global_orient_mat)
    position_data = torch.einsum('ij,tkj->tki', rotation_matrix, torch.from_numpy(joints).float())
    position_val_data = position_data[1:] - position_data[:-1]

    root_idx = 0
    # put on floor and put root on origin for the first frame
    ori = deepcopy(position_data[0, root_idx]) # first frame root position
    y_min = torch.min(position_data[:, :, 1])
    # ori[1] = y_min
    # position_data = position_data - ori
    velocities_root = position_data[1:, root_idx, :] - position_data[:-1, root_idx, :]
    position_data_cp = deepcopy(position_data)
    position_data[:,:,0] -= position_data_cp[:,0:1,0]
    position_data[:,:,2] -= position_data_cp[:,0:1,2]

    # vis_3d_g1([position_data.numpy()[:, 1:]], None, ['video.mp4'], fps=30)
    T, njoint, _ = position_data.shape
    final_x = torch.zeros((T, 2 + 6 + njoint * 3 + njoint * 3))
    final_x[1:, 0] = velocities_root[:, 0]
    final_x[1:, 1] = velocities_root[:, 2]
    final_x[:, 2:2+6] = matrix_to_rotation_6d(global_orient_mat)
    final_x[:, 8:8+njoint*3] = position_data.flatten(1, 2)
    final_x[1:, 8+njoint*3:8+njoint*6] = position_val_data.flatten(1, 2) # T, 140
    final_x = torch.concat([final_x, torch.from_numpy(dof).float()], dim=-1)
    # if final_x.shape[0] > 200:
    #     import ipdb; ipdb.set_trace()
    # dof, rot_quat, transl = extract_g1_component(final_x)
    # joblib.dump(dict(dof=dof, rot_quat=rot_quat, transl=transl), 'data.pkl')
    return final_x # 217

# python -m data.motionmillion.tools.process --save_root "data/motionmillion/final_data"

MUJOCO_DOF_COLUMNS = [
    'left_hip_pitch_joint_dof', 'left_hip_roll_joint_dof', 'left_hip_yaw_joint_dof',
    'left_knee_joint_dof', 'left_ankle_pitch_joint_dof', 'left_ankle_roll_joint_dof',
    'right_hip_pitch_joint_dof', 'right_hip_roll_joint_dof', 'right_hip_yaw_joint_dof',
    'right_knee_joint_dof', 'right_ankle_pitch_joint_dof', 'right_ankle_roll_joint_dof',
    'waist_yaw_joint_dof', 'waist_roll_joint_dof', 'waist_pitch_joint_dof',
    'left_shoulder_pitch_joint_dof', 'left_shoulder_roll_joint_dof', 'left_shoulder_yaw_joint_dof',
    'left_elbow_joint_dof', 'left_wrist_roll_joint_dof', 'left_wrist_pitch_joint_dof',
    'left_wrist_yaw_joint_dof',
    'right_shoulder_pitch_joint_dof', 'right_shoulder_roll_joint_dof', 'right_shoulder_yaw_joint_dof',
    'right_elbow_joint_dof', 'right_wrist_roll_joint_dof', 'right_wrist_pitch_joint_dof',
    'right_wrist_yaw_joint_dof',
]

# MuJoCo DFS DOF index -> IsaacLab BFS DOF index
MUJOCO_TO_ISAAC = [0, 3, 6, 9, 13, 17, 1, 4, 7, 10, 14, 18, 2, 5, 8, 11,
                   15, 19, 21, 23, 25, 27, 12, 16, 20, 22, 24, 26, 28]

# pinocchio body index -> IsaacLab body index (for 30 bodies, including root)
body_mapping = [2, 4, 20, 34, 6, 22, 36, 8, 24, 38, 10, 26, 46, 62, 12, 28, 48, 64, 14, 30, 50, 66, 52, 68, 54, 70, 56, 72, 58, 74]

URDF_PATH = 'tools/GMR/assets/unitree_g1/g1_custom_collision_29dof.urdf'

# --- Forward kinematics ---
def compute_fk(model, dof, root_pos, root_quat_wxyz):
    """Compute 30 body world positions and orientations via Pinocchio forward kinematics.
    Args:
        model: Pinocchio model with free-flyer base
        dof: (T, 29) DOF angles, radians, Mujoco order
        root_pos: (T, 3) root position in meters
        root_quat_wxyz: (T, 4) root quaternion wxyz
    Returns:
        body_pos_w:  (T, 30, 3) body positions world frame
        body_quat_w: (T, 30, 4) body orientations world frame, wxyz
    """
    data = model.createData()
    T = dof.shape[0]

    body_pos_w = np.zeros((T, 30, 3))
    body_rot_w = np.zeros((T, 30, 3, 3))
    for t in range(T):
        q = np.zeros(model.nq)
        q[0:3] = root_pos[t]
        # Pinocchio quaternion in xyzw order
        q[3:7] = root_quat_wxyz[t, [1, 2, 3, 0]]
        q[7:] = dof[t]
        # q[7:] = mapped_joint[t]

        pin.forwardKinematics(model, data, q)
        pin.updateFramePlacements(model, data)

        for lab_idx, pino_frame_idx in enumerate(body_mapping):
            body_pos_w[t, lab_idx] = data.oMf[pino_frame_idx].translation
            body_rot_w[t, lab_idx] = data.oMf[pino_frame_idx].rotation

    # Rotation matrices -> wxyz quaternions
    rot_flat = R.from_matrix(body_rot_w.reshape(-1, 3, 3))
    quat_xyzw = rot_flat.as_quat().reshape(T, 30, 4)
    body_quat_w = quat_xyzw[..., [3, 0, 1, 2]]  # xyzw -> wxyz

    return body_pos_w, body_quat_w


# --- Main function ---

def csv_to_gmr_npz(csv_path: str, source_fps: int = 120, target_fps: int = 30):
    """Load MuJoCo CSV, downsample, compute FK, save as GMR-compatible npz.
    Args:
        csv_path: path to input CSV file (at source_fps)
        save_path: path to output npz file
        source_fps: CSV frame rate (default 120)
        target_fps: output frame rate after downsampling (default 30)
    """
    # Step 1: Read CSV and immediately downsample by stride
    assert source_fps % target_fps == 0, \
        f'source_fps ({source_fps}) must be a multiple of target_fps ({target_fps})'
    stride = source_fps // target_fps

    df = pd.read_csv(csv_path)
    T_src = len(df)
    df = df.iloc[::stride].reset_index(drop=True)
    T_tgt = len(df)
    # Step 2: Unit conversions
    root_pos = df[['root_translateX', 'root_translateY', 'root_translateZ']].values / 100.0  # cm -> m
    dof_mujoco_rad = np.deg2rad(df[MUJOCO_DOF_COLUMNS].values)  # deg -> rad, (T_tgt, 29)

    # Step 3: Root euler (degrees, extrinsic XYZ) -> wxyz quaternion
    euler_deg = df[['root_rotateX', 'root_rotateY', 'root_rotateZ']].values
    rot = R.from_euler('xyz', euler_deg, degrees=True)
    quat_xyzw = rot.as_quat()
    root_quat_wxyz = quat_xyzw[:, [3, 0, 1, 2]].astype(np.float64)

    # Step 4: DOF reorder MuJoCo -> IsaacLab
    dof_isaac = np.zeros_like(dof_mujoco_rad)
    for mj_idx, isaac_idx in enumerate(MUJOCO_TO_ISAAC):
        dof_isaac[:, isaac_idx] = dof_mujoco_rad[:, mj_idx]

    # Step 5: Pinocchio FK to get 30 body positions + orientations (only on downsampled frames)
    model = pin.buildModelFromUrdf(URDF_PATH, pin.JointModelFreeFlyer())
    body_pos_w, body_quat_w = compute_fk(model, dof_mujoco_rad, root_pos, root_quat_wxyz)
    body_quat_w = body_quat_w.astype(np.float64)

    joint_pos=dof_isaac.astype(np.float64)
    body_pos_w=body_pos_w.astype(np.float64)
    return joint_pos, body_pos_w, body_quat_w

def func(line):
    data_path = line.strip()
    joint_pos, body_pos_w, body_quat_w = csv_to_gmr_npz(data_path)
    g1_motion = get_g1_motion(joint_pos, body_pos_w, body_quat_w)   
    # import ipdb; ipdb.set_trace()
    data_path = data_path.replace('.csv', '.npy').replace('/mnt/shenzhen2cephfs/capybarali/seed/seed/g1/csv/', '')
    save_path = args.save_root + '/' + data_path
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    np.save(save_path, g1_motion)

if __name__ == '__main__':
    with open('/mnt/shenzhen2cephfs/capybarali/seed/seed_g1_csv_path.txt', 'r') as f:
        paths = f.readlines()
    Parallel(n_jobs=64)(delayed(func)(line) for line in tqdm(paths[args.start_idx::args.interval]))

	import sys
	sys.path.append('/mnt/shenzhen2cephfs/capybarali/codes/humanoid')

	import torch, yaml, os
	from tqdm import tqdm
	from src.utils.rotation_conversions import quaternion_to_matrix, matrix_to_rotation_6d, matrix_to_axis_angle, rotation_6d_to_matrix, matrix_to_quaternion, axis_angle_to_matrix
	import numpy as np
	from argparse import ArgumentParser
	import joblib
	from data.vis import vis_3d_motion
	from data.vis_g1 import vis_3d_g1
	from copy import deepcopy
	from scipy.spatial.transform import Rotation as R
	from scipy.spatial.transform import Slerp
	from scipy import interpolate
	from joblib import Parallel, delayed
	import argparse
	import numpy as np
	import pandas as pd
	from scipy.spatial.transform import Rotation as R
	import pinocchio as pin

	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	parser = ArgumentParser(description="Launch MoCap processing")
	parser.add_argument('--save_root', type=str, default="data/seed_g1_217")
	parser.add_argument('--start_idx', type=int, default=0)
	parser.add_argument('--interval', type=int, default=1)
	args = parser.parse_args()

	os.makedirs(args.save_root, exist_ok=True)

	def extract_g1_component(x):
	vel_xy = x[:, :2]
	dof = x[:, -29:]
	root_rot_mat = rotation_6d_to_matrix(x[:, 2:8])
	trans_xyz = torch.cat([
	torch.cumsum(vel_xy[:, :1], dim=0), # new_x = cumsum(Δnew_x)
	x[:, 9:10], # new_y = height (direct)
	torch.cumsum(vel_xy[:, 1:], dim=0), # new_z = cumsum(Δnew_z)
	], dim=1)

	rot_mat = torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]]).float()
	global_orient_mat = root_rot_mat.squeeze(1).float()
	global_orient_mat = torch.einsum('ij,tjk->tik', rot_mat, global_orient_mat)
	rot_quat = matrix_to_quaternion(global_orient_mat) # (T, 4) wxyz order

	transl = trans_xyz.float()
	transl = torch.einsum('ij,tj->ti', rot_mat, transl)

	return dof, rot_quat, transl

	def get_g1_motion(joint_pos, body_pos_w, body_quat_w):
	dof = joint_pos # T, 29
	root_ori = body_quat_w[:, 0] # T, 4, wxyz
	joints = body_pos_w # T, J, 3
	# import ipdb; ipdb.set_trace()
	# 插值
	rotation_matrix = torch.tensor([[1.0, 0, 0], [0, 0, -1], [0, 1, 0]]).inverse()

	global_orient_mat = quaternion_to_matrix(torch.from_numpy(root_ori)).float()
	global_orient_mat = torch.einsum('ij,tjk->tik', rotation_matrix, global_orient_mat)
	global_orient = matrix_to_axis_angle(global_orient_mat)
	position_data = torch.einsum('ij,tkj->tki', rotation_matrix, torch.from_numpy(joints).float())
	position_val_data = position_data[1:] - position_data[:-1]

	root_idx = 0
	# put on floor and put root on origin for the first frame
	ori = deepcopy(position_data[0, root_idx]) # first frame root position
	y_min = torch.min(position_data[:, :, 1])
	# ori[1] = y_min
	# position_data = position_data - ori
	velocities_root = position_data[1:, root_idx, :] - position_data[:-1, root_idx, :]
	position_data_cp = deepcopy(position_data)
	position_data[:,:,0] -= position_data_cp[:,0:1,0]
	position_data[:,:,2] -= position_data_cp[:,0:1,2]

	# vis_3d_g1([position_data.numpy()[:, 1:]], None, ['video.mp4'], fps=30)
	T, njoint, _ = position_data.shape
	final_x = torch.zeros((T, 2 + 6 + njoint * 3 + njoint * 3))
	final_x[1:, 0] = velocities_root[:, 0]
	final_x[1:, 1] = velocities_root[:, 2]
	final_x[:, 2:2+6] = matrix_to_rotation_6d(global_orient_mat)
	final_x[:, 8:8+njoint*3] = position_data.flatten(1, 2)
	final_x[1:, 8+njoint3:8+njoint6] = position_val_data.flatten(1, 2) # T, 140
	final_x = torch.concat([final_x, torch.from_numpy(dof).float()], dim=-1)
	# if final_x.shape[0] > 200:
	# import ipdb; ipdb.set_trace()
	# dof, rot_quat, transl = extract_g1_component(final_x)
	# joblib.dump(dict(dof=dof, rot_quat=rot_quat, transl=transl), 'data.pkl')
	return final_x # 217

	# python -m data.motionmillion.tools.process --save_root "data/motionmillion/final_data"

	MUJOCO_DOF_COLUMNS = [
	'left_hip_pitch_joint_dof', 'left_hip_roll_joint_dof', 'left_hip_yaw_joint_dof',
	'left_knee_joint_dof', 'left_ankle_pitch_joint_dof', 'left_ankle_roll_joint_dof',
	'right_hip_pitch_joint_dof', 'right_hip_roll_joint_dof', 'right_hip_yaw_joint_dof',
	'right_knee_joint_dof', 'right_ankle_pitch_joint_dof', 'right_ankle_roll_joint_dof',
	'waist_yaw_joint_dof', 'waist_roll_joint_dof', 'waist_pitch_joint_dof',
	'left_shoulder_pitch_joint_dof', 'left_shoulder_roll_joint_dof', 'left_shoulder_yaw_joint_dof',
	'left_elbow_joint_dof', 'left_wrist_roll_joint_dof', 'left_wrist_pitch_joint_dof',
	'left_wrist_yaw_joint_dof',
	'right_shoulder_pitch_joint_dof', 'right_shoulder_roll_joint_dof', 'right_shoulder_yaw_joint_dof',
	'right_elbow_joint_dof', 'right_wrist_roll_joint_dof', 'right_wrist_pitch_joint_dof',
	'right_wrist_yaw_joint_dof',
	]

	# MuJoCo DFS DOF index -> IsaacLab BFS DOF index
	MUJOCO_TO_ISAAC = [0, 3, 6, 9, 13, 17, 1, 4, 7, 10, 14, 18, 2, 5, 8, 11,
	15, 19, 21, 23, 25, 27, 12, 16, 20, 22, 24, 26, 28]

	# pinocchio body index -> IsaacLab body index (for 30 bodies, including root)
	body_mapping = [2, 4, 20, 34, 6, 22, 36, 8, 24, 38, 10, 26, 46, 62, 12, 28, 48, 64, 14, 30, 50, 66, 52, 68, 54, 70, 56, 72, 58, 74]

	URDF_PATH = 'tools/GMR/assets/unitree_g1/g1_custom_collision_29dof.urdf'

	# --- Forward kinematics ---
	def compute_fk(model, dof, root_pos, root_quat_wxyz):
	"""Compute 30 body world positions and orientations via Pinocchio forward kinematics.

	Args:
	model: Pinocchio model with free-flyer base
	dof: (T, 29) DOF angles, radians, Mujoco order
	root_pos: (T, 3) root position in meters
	root_quat_wxyz: (T, 4) root quaternion wxyz

	Returns:
	body_pos_w: (T, 30, 3) body positions world frame
	body_quat_w: (T, 30, 4) body orientations world frame, wxyz
	"""
	data = model.createData()
	T = dof.shape[0]

	body_pos_w = np.zeros((T, 30, 3))
	body_rot_w = np.zeros((T, 30, 3, 3))
	for t in range(T):
	q = np.zeros(model.nq)
	q[0:3] = root_pos[t]
	# Pinocchio quaternion in xyzw order
	q[3:7] = root_quat_wxyz[t, [1, 2, 3, 0]]
	q[7:] = dof[t]
	# q[7:] = mapped_joint[t]

	pin.forwardKinematics(model, data, q)
	pin.updateFramePlacements(model, data)

	for lab_idx, pino_frame_idx in enumerate(body_mapping):
	body_pos_w[t, lab_idx] = data.oMf[pino_frame_idx].translation
	body_rot_w[t, lab_idx] = data.oMf[pino_frame_idx].rotation

	# Rotation matrices -> wxyz quaternions
	rot_flat = R.from_matrix(body_rot_w.reshape(-1, 3, 3))
	quat_xyzw = rot_flat.as_quat().reshape(T, 30, 4)
	body_quat_w = quat_xyzw[..., [3, 0, 1, 2]] # xyzw -> wxyz

	return body_pos_w, body_quat_w


	# --- Main function ---

	def csv_to_gmr_npz(csv_path: str, source_fps: int = 120, target_fps: int = 30):
	"""Load MuJoCo CSV, downsample, compute FK, save as GMR-compatible npz.

	Args:
	csv_path: path to input CSV file (at source_fps)
	save_path: path to output npz file
	source_fps: CSV frame rate (default 120)
	target_fps: output frame rate after downsampling (default 30)
	"""
	# Step 1: Read CSV and immediately downsample by stride
	assert source_fps % target_fps == 0, \
	f'source_fps ({source_fps}) must be a multiple of target_fps ({target_fps})'
	stride = source_fps // target_fps

	df = pd.read_csv(csv_path)
	T_src = len(df)
	df = df.iloc[::stride].reset_index(drop=True)
	T_tgt = len(df)
	# Step 2: Unit conversions
	root_pos = df[['root_translateX', 'root_translateY', 'root_translateZ']].values / 100.0 # cm -> m
	dof_mujoco_rad = np.deg2rad(df[MUJOCO_DOF_COLUMNS].values) # deg -> rad, (T_tgt, 29)

	# Step 3: Root euler (degrees, extrinsic XYZ) -> wxyz quaternion
	euler_deg = df[['root_rotateX', 'root_rotateY', 'root_rotateZ']].values
	rot = R.from_euler('xyz', euler_deg, degrees=True)
	quat_xyzw = rot.as_quat()
	root_quat_wxyz = quat_xyzw[:, [3, 0, 1, 2]].astype(np.float64)

	# Step 4: DOF reorder MuJoCo -> IsaacLab
	dof_isaac = np.zeros_like(dof_mujoco_rad)
	for mj_idx, isaac_idx in enumerate(MUJOCO_TO_ISAAC):
	dof_isaac[:, isaac_idx] = dof_mujoco_rad[:, mj_idx]

	# Step 5: Pinocchio FK to get 30 body positions + orientations (only on downsampled frames)
	model = pin.buildModelFromUrdf(URDF_PATH, pin.JointModelFreeFlyer())
	body_pos_w, body_quat_w = compute_fk(model, dof_mujoco_rad, root_pos, root_quat_wxyz)
	body_quat_w = body_quat_w.astype(np.float64)

	joint_pos=dof_isaac.astype(np.float64)
	body_pos_w=body_pos_w.astype(np.float64)
	return joint_pos, body_pos_w, body_quat_w

	def func(line):
	data_path = line.strip()
	joint_pos, body_pos_w, body_quat_w = csv_to_gmr_npz(data_path)
	g1_motion = get_g1_motion(joint_pos, body_pos_w, body_quat_w)
	# import ipdb; ipdb.set_trace()
	data_path = data_path.replace('.csv', '.npy').replace('/mnt/shenzhen2cephfs/capybarali/seed/seed/g1/csv/', '')
	save_path = args.save_root + '/' + data_path
	os.makedirs(os.path.dirname(save_path), exist_ok=True)
	np.save(save_path, g1_motion)

	if __name__ == '__main__':
	with open('/mnt/shenzhen2cephfs/capybarali/seed/seed_g1_csv_path.txt', 'r') as f:
	paths = f.readlines()
	Parallel(n_jobs=64)(delayed(func)(line) for line in tqdm(paths[args.start_idx::args.interval]))