import os from os.path import join as pjoin import glob import json import logging import torch import numpy as np from pybullet_utils import transformations from rsl_rl.utils import motion_util _EPS = np.finfo(float).eps * 4.0 def quaternion_slerp(q0, q1, fraction, spin=0, shortestpath=True): """Batch quaternion spherical linear interpolation.""" out = torch.zeros_like(q0) zero_mask = torch.isclose(fraction, torch.zeros_like(fraction)).squeeze() ones_mask = torch.isclose(fraction, torch.ones_like(fraction)).squeeze() out[zero_mask] = q0[zero_mask] out[ones_mask] = q1[ones_mask] d = torch.sum(q0 * q1, dim=-1, keepdim=True) dist_mask = (torch.abs(torch.abs(d) - 1.0) < _EPS).squeeze() out[dist_mask] = q0[dist_mask] if shortestpath: d_old = torch.clone(d) d = torch.where(d_old < 0, -d, d) q1 = torch.where(d_old < 0, -q1, q1) angle = torch.acos(d) + spin * torch.pi angle_mask = (torch.abs(angle) < _EPS).squeeze() out[angle_mask] = q0[angle_mask] final_mask = torch.logical_or(zero_mask, ones_mask) final_mask = torch.logical_or(final_mask, dist_mask) final_mask = torch.logical_or(final_mask, angle_mask) final_mask = torch.logical_not(final_mask) isin = 1.0 / angle q0 *= torch.sin((1.0 - fraction) * angle) * isin q1 *= torch.sin(fraction * angle) * isin q0 += q1 out[final_mask] = q0[final_mask] return out class G1_AMPLoader: def __init__( self, device, time_between_frames, motion_files, preload_transitions=False, num_preload_transitions=1000000, num_frames=5, ): """Expert dataset provides AMP observations from Dog mocap dataset. time_between_frames: Amount of time in seconds between transition. """ self.device = device self.time_between_frames = time_between_frames self.num_frames = num_frames # Values to store for each trajectory. self.trajectories = [] self.trajectories_full = [] self.trajectory_names = [] self.trajectory_idxs = [] self.trajectory_lens = [] # Traj length in seconds. self.trajectory_weights = [] self.trajectory_frame_durations = [] self.trajectory_num_frames = [] self.motion_dir = motion_files # import ipdb; ipdb.set_trace() for i, motion_file in enumerate(os.listdir(motion_files)): self.trajectory_names.append(motion_file) motion_path = pjoin(motion_files, motion_file) motion_data = np.load(motion_path, allow_pickle=True) motion_data_processed = np.zeros((motion_data.shape[0],36)) for f_i in range(motion_data.shape[0]): motion_data_processed[f_i, :3] = motion_data[f_i, :3] # base pos motion_data_processed[f_i, 3:7] = motion_data[f_i, 3:7] # base quat (wxyz) motion_data_processed[f_i, 7:35] = motion_data[f_i, 7:35] # base vel ''' NOTE The order of motion_data_processed is base pos 0:3, base quat 3:7, wxyz dof pos 7:36, (mujoco joint order) ''' self.trajectories.append(torch.tensor( motion_data_processed[:, 7:], dtype=torch.float32, device=self.device )) self.trajectories_full.append(torch.tensor( motion_data_processed, dtype=torch.float32, device=self.device )) self.trajectory_idxs.append(i) self.trajectory_weights.append(1 / len(os.listdir(motion_files))) frame_duration = 1 / 50 self.trajectory_frame_durations.append(frame_duration) traj_len = (motion_data_processed.shape[0] - 1) * frame_duration # seconds self.trajectory_lens.append(traj_len) self.trajectory_num_frames.append(float(motion_data_processed.shape[0])) print(f"Loaded {traj_len}s. motion from {motion_file}.") # Trajectory weights are used to sample some trajectories more than others. self.trajectory_weights = np.array(self.trajectory_weights) / np.sum(self.trajectory_weights) self.trajectory_frame_durations = np.array(self.trajectory_frame_durations) self.trajectory_lens = np.array(self.trajectory_lens) self.trajectory_num_frames = np.array(self.trajectory_num_frames) # Preload transitions. self.preload_transitions = preload_transitions if self.preload_transitions: print(f'Preloading {num_preload_transitions} transitions') traj_idxs = self.weighted_traj_idx_sample_batch(num_preload_transitions) times = self.traj_time_sample_batch(traj_idxs) self.preloaded_s_prior = self.get_full_frame_at_time_batch(traj_idxs, times - self.time_between_frames) self.preloaded_s = self.get_full_frame_at_time_batch(traj_idxs, times) self.preloaded_s_next = self.get_full_frame_at_time_batch(traj_idxs, times + self.time_between_frames) print(f'Finished preloading') # 预加载多帧数据 self.preloaded_frames = [] for i in range(self.num_frames): frame_time = times + (i - (self.num_frames - 2)) * self.time_between_frames full_frame = self.get_full_frame_at_time_batch(traj_idxs, frame_time) # 预处理：提前提取并连接需要的列（7:26 和 29:33），避免每次生成时重复切片 processed_frame = torch.cat([ full_frame[:, 7:26], full_frame[:, 29:33] ], dim=-1) self.preloaded_frames.append(processed_frame) print(f'Finished preloading multiple frames') self.all_trajectories_full = torch.vstack(self.trajectories_full) def weighted_traj_idx_sample(self): """Get traj idx via weighted sampling.""" return np.random.choice( self.trajectory_idxs, p=self.trajectory_weights) def weighted_traj_idx_sample_batch(self, size): """Batch sample traj idxs.""" return np.random.choice( self.trajectory_idxs, size=size, p=self.trajectory_weights, replace=True) def traj_time_sample(self, traj_idx): """Sample random time for traj.""" subst = self.time_between_frames + self.trajectory_frame_durations[traj_idx] return max( 0, (self.trajectory_lens[traj_idx] * np.random.uniform() - subst)) def traj_time_sample_batch(self, traj_idxs): """Sample random time for multiple trajectories.""" subst = self.time_between_frames + self.trajectory_frame_durations[traj_idxs] time_samples = self.trajectory_lens[traj_idxs] * np.random.uniform(size=len(traj_idxs)) - subst return np.maximum(np.zeros_like(time_samples), time_samples) def slerp(self, val0, val1, blend): return (1.0 - blend) * val0 + blend * val1 def get_trajectory(self, traj_idx): """Returns trajectory of AMP observations.""" return self.trajectories_full[traj_idx] def get_frame_at_time(self, traj_idx, time): """Returns frame for the given trajectory at the specified time.""" p = float(time) / self.trajectory_lens[traj_idx] n = self.trajectories[traj_idx].shape[0] idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) frame_start = self.trajectories[traj_idx][idx_low] frame_end = self.trajectories[traj_idx][idx_high] blend = p * n - idx_low return self.slerp(frame_start, frame_end, blend) def get_frame_at_time_batch(self, traj_idxs, times): """Returns frame for the given trajectory at the specified time.""" p = times / self.trajectory_lens[traj_idxs] n = self.trajectory_num_frames[traj_idxs] idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) all_frame_starts = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) all_frame_ends = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) for traj_idx in set(traj_idxs): trajectory = self.trajectories[traj_idx] traj_mask = traj_idxs == traj_idx all_frame_starts[traj_mask] = trajectory[idx_low[traj_mask]] all_frame_ends[traj_mask] = trajectory[idx_high[traj_mask]] blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) return self.slerp(all_frame_starts, all_frame_ends, blend) def get_full_frame_at_time(self, traj_idx, time): """Returns full frame for the given trajectory at the specified time.""" p = float(time) / self.trajectory_lens[traj_idx] n = self.trajectories_full[traj_idx].shape[0] idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) frame_start = self.trajectories_full[traj_idx][idx_low] frame_end = self.trajectories_full[traj_idx][idx_high] blend = p * n - idx_low print(idx_low, idx_high) return self.blend_frame_pose(frame_start, frame_end, blend) def get_full_frame_at_time_batch(self, traj_idxs, times): p = times / self.trajectory_lens[traj_idxs] n = self.trajectory_num_frames[traj_idxs] idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) all_frame_pos_starts = torch.zeros(len(traj_idxs), 3, device=self.device) all_frame_pos_ends = torch.zeros(len(traj_idxs), 3, device=self.device) all_frame_rot_starts = torch.zeros(len(traj_idxs), 4, device=self.device) all_frame_rot_ends = torch.zeros(len(traj_idxs), 4, device=self.device) all_frame_amp_starts = torch.zeros(len(traj_idxs), 29, device=self.device) all_frame_amp_ends = torch.zeros(len(traj_idxs), 29, device=self.device) for traj_idx in set(traj_idxs): trajectory = self.trajectories_full[traj_idx] traj_mask = traj_idxs == traj_idx all_frame_pos_starts[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_low[traj_mask]]) all_frame_pos_ends[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_high[traj_mask]]) all_frame_rot_starts[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_low[traj_mask]]) all_frame_rot_ends[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_high[traj_mask]]) all_frame_amp_starts[traj_mask] = trajectory[idx_low[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 all_frame_amp_ends[traj_mask] = trajectory[idx_high[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) pos_blend = self.slerp(all_frame_pos_starts, all_frame_pos_ends, blend) rot_blend = quaternion_slerp(all_frame_rot_starts, all_frame_rot_ends, blend) amp_blend = self.slerp(all_frame_amp_starts, all_frame_amp_ends, blend) return torch.cat([pos_blend, rot_blend, amp_blend], dim=-1) def get_frame(self): """Returns random frame.""" traj_idx = self.weighted_traj_idx_sample() sampled_time = self.traj_time_sample(traj_idx) return self.get_frame_at_time(traj_idx, sampled_time) def get_full_frame(self): """Returns random full frame.""" traj_idx = self.weighted_traj_idx_sample() sampled_time = self.traj_time_sample(traj_idx) return self.get_full_frame_at_time(traj_idx, sampled_time) def get_full_frame_batch(self, num_frames): if self.preload_transitions: idxs = np.random.choice( self.preloaded_s.shape[0], size=num_frames) return self.preloaded_s[idxs] else: traj_idxs = self.weighted_traj_idx_sample_batch(num_frames) times = self.traj_time_sample_batch(traj_idxs) return self.get_full_frame_at_time_batch(traj_idxs, times) def blend_frame_pose(self, frame0, frame1, blend): """Linearly interpolate between two frames, including orientation. Args: frame0: First frame to be blended corresponds to (blend = 0). frame1: Second frame to be blended corresponds to (blend = 1). blend: Float between [0, 1], specifying the interpolation between the two frames. Returns: An interpolation of the two frames. """ root_pos0, root_pos1 = G1_AMPLoader.get_root_pos(frame0), G1_AMPLoader.get_root_pos(frame1) root_rot0, root_rot1 = G1_AMPLoader.get_root_rot(frame0), G1_AMPLoader.get_root_rot(frame1) joints0, joints1 = G1_AMPLoader.get_joint_pose(frame0), G1_AMPLoader.get_joint_pose(frame1) # tar_toe_pos_0, tar_toe_pos_1 = G1_AMPLoader.get_tar_toe_pos_local(frame0), G1_AMPLoader.get_tar_toe_pos_local(frame1) linear_vel_0, linear_vel_1 = G1_AMPLoader.get_linear_vel(frame0), G1_AMPLoader.get_linear_vel(frame1) angular_vel_0, angular_vel_1 = G1_AMPLoader.get_angular_vel(frame0), G1_AMPLoader.get_angular_vel(frame1) joint_vel_0, joint_vel_1 = G1_AMPLoader.get_joint_vel(frame0), G1_AMPLoader.get_joint_vel(frame1) blend_root_pos = self.slerp(root_pos0, root_pos1, blend) blend_root_rot = transformations.quaternion_slerp(root_rot0.cpu().numpy(), root_rot1.cpu().numpy(), blend) blend_root_rot = torch.tensor(motion_util.standardize_quaternion(blend_root_rot),dtype=torch.float32, device=self.device) blend_joints = self.slerp(joints0, joints1, blend) # blend_tar_toe_pos = self.slerp(tar_toe_pos_0, tar_toe_pos_1, blend) blend_linear_vel = self.slerp(linear_vel_0, linear_vel_1, blend) blend_angular_vel = self.slerp(angular_vel_0, angular_vel_1, blend) blend_joints_vel = self.slerp(joint_vel_0, joint_vel_1, blend) # return # torch.cat([ # blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints, blend_joints_vel]) return torch.cat([blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints]) def feed_forward_generator_23dof_multi(self, num_mini_batch, mini_batch_size): """Generates a batch of AMP transitions.""" # import ipdb; ipdb.set_trace() for _ in range(num_mini_batch): if self.preload_transitions: idxs = np.random.choice(self.preloaded_s.shape[0], size=mini_batch_size) frames = [] for i in range(self.num_frames): # 数据已在预加载时预处理，直接索引即可 s = self.preloaded_frames[i][idxs] frames.append(s) else: NotImplementedError('preload transition') yield torch.stack(frames, dim=1) # [batch, num_frames, 16] def quaternion_to_euler_array(self, quat): # Ensure quaternion is in the correct format [x, y, z, w] x, y, z, w =quat # Roll (x-axis rotation) t0 = +2.0 * (w * x + y * z) t1 = +1.0 - 2.0 * (x * x + y * y) roll_x = np.arctan2(t0, t1) # Pitch (y-axis rotation) t2 = +2.0 * (w * y - z * x) t2 = np.clip(t2, -1.0, 1.0) pitch_y = np.arcsin(t2) # Yaw (z-axis rotation) t3 = +2.0 * (w * z + x * y) t4 = +1.0 - 2.0 * (y * y + z * z) yaw_z = np.arctan2(t3, t4) # Returns roll, pitch, yaw in a NumPy array in radians return np.array([roll_x, pitch_y, yaw_z]) def euler_to_quaternion(self, root_rot): roll, pitch, yaw = root_rot[0], root_rot[1], root_rot[2] cy = np.cos(yaw * 0.5) sy = np.sin(yaw * 0.5) cp = np.cos(pitch * 0.5) sp = np.sin(pitch * 0.5) cr = np.cos(roll * 0.5) sr = np.sin(roll * 0.5) qw = cy * cp * cr + sy * sp * sr qx = cy * cp * sr - sy * sp * cr qy = sy * cp * sr + cy * sp * cr qz = sy * cp * cr - cy * sp * sr return np.array([qx, qy, qz, qw]) @property def observation_dim(self): """Size of AMP observations.""" return self.trajectories[0].shape[1] + 1 @property def num_motions(self): return len(self.trajectory_names) @staticmethod def get_root_pos(pose): return pose[0:3] @staticmethod def get_root_pos_batch(poses): return poses[:, 0:3] @staticmethod def get_root_rot(pose): return pose[3:7] @staticmethod def get_root_rot_batch(poses): return poses[:, 3:7] @staticmethod def get_joint_pose_batch_12dof(poses): return poses[:, 13:25] @staticmethod def get_tar_toe_pos_local(pose): return pose[G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] @staticmethod def get_tar_toe_pos_local_batch(poses): return poses[:, G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX]