| import enum |
| import gc |
| import glob |
| import os |
| import os.path as osp |
| from pathlib import Path |
| import random |
| import re |
| import resource |
|
|
| import easydict |
| import joblib |
| from loguru import logger |
| import numpy as np |
| from rich import progress |
| from scipy.spatial import transform |
| import torch |
| import torch.multiprocessing as mp |
|
|
| from gear_sonic.isaac_utils import rotations |
| from gear_sonic.trl.utils import common |
| from gear_sonic.utils.motion_lib import skeleton |
|
|
|
|
| class FixHeightMode(enum.Enum): |
| no_fix = 0 |
| full_fix = 1 |
| ankle_fix = 2 |
|
|
|
|
| class MotionlibMode(enum.Enum): |
| file = 1 |
| directory = 2 |
|
|
|
|
| def to_torch(tensor): |
| if torch.is_tensor(tensor): |
| return tensor |
| else: |
| return torch.from_numpy(tensor) |
|
|
|
|
| def is_navigation_motion(motion_key): |
| return ( |
| motion_key.startswith("2025") |
| or motion_key.startswith("walking_2025") |
| or motion_key.startswith("running_2025") |
| or motion_key.startswith("slow_walk_2025") |
| ) |
|
|
|
|
| def interpolate_translation_data( |
| data, |
| source_fps, |
| target_fps, |
| num_frames, |
| max_num_objects=1, |
| pad_value=0.0, |
| ): |
| """Interpolate translation-like data (e.g., root_pos, contact_points) to target FPS. |
| |
| Args: |
| data: Tensor of shape (T, N, D) where T=frames, N=num_objects, D=dims (e.g., 3 for pos) |
| source_fps: Original frame rate |
| target_fps: Target frame rate |
| num_frames: Target number of frames |
| max_num_objects: Maximum number of objects to pad to |
| pad_value: Value to use for padding |
| |
| Returns: |
| Interpolated tensor of shape (num_frames, max_num_objects, D) |
| """ |
| from gear_sonic.trl.utils import math |
|
|
| data = torch.tensor(data).float() if not torch.is_tensor(data) else data.float() |
| N_objects = data.shape[1] |
| D = data.shape[2] |
|
|
| |
| if source_fps != target_fps: |
| |
| data_flat = data.reshape(data.shape[0], -1) |
| data_interp = math.interpolate_pose( |
| data_flat, |
| source_fps=source_fps, |
| target_fps=target_fps, |
| device=data.device, |
| interpolation_type="linear", |
| ) |
| data = data_interp.reshape(-1, N_objects, D) |
|
|
| |
| if data.shape[0] > num_frames: |
| data = data[:num_frames] |
| elif data.shape[0] < num_frames: |
| padding = data[-1:].repeat(num_frames - data.shape[0], 1, 1) |
| data = torch.cat([data, padding], dim=0) |
|
|
| |
| if data.shape[1] < max_num_objects: |
| padding = torch.full( |
| (data.shape[0], max_num_objects - data.shape[1], D), |
| pad_value, |
| dtype=data.dtype, |
| device=data.device, |
| ) |
| data = torch.cat([data, padding], dim=1) |
| else: |
| data = data[:, :max_num_objects] |
|
|
| return data |
|
|
|
|
| def interpolate_quaternion_data( |
| data, |
| source_fps, |
| target_fps, |
| num_frames, |
| max_num_objects=1, |
| ): |
| """Interpolate quaternion data (e.g., root_quat) to target FPS using slerp. |
| |
| Args: |
| data: Tensor of shape (T, N, 4) where T=frames, N=num_objects |
| source_fps: Original frame rate |
| target_fps: Target frame rate |
| num_frames: Target number of frames |
| max_num_objects: Maximum number of objects to pad to |
| |
| Returns: |
| Interpolated tensor of shape (num_frames, max_num_objects, 4) |
| """ |
| from gear_sonic.trl.utils import math |
|
|
| data = torch.tensor(data).float() if not torch.is_tensor(data) else data.float() |
| N_objects = data.shape[1] |
|
|
| |
| if source_fps != target_fps: |
| |
| data_flat = data.reshape(data.shape[0], -1) |
| data_interp = math.interpolate_pose( |
| data_flat, |
| source_fps=source_fps, |
| target_fps=target_fps, |
| device=data.device, |
| interpolation_type="slerp", |
| rot_type="quat", |
| ) |
| data = data_interp.reshape(-1, N_objects, 4) |
|
|
| |
| if data.shape[0] > num_frames: |
| data = data[:num_frames] |
| elif data.shape[0] < num_frames: |
| padding = data[-1:].repeat(num_frames - data.shape[0], 1, 1) |
| data = torch.cat([data, padding], dim=0) |
|
|
| |
| if data.shape[1] < max_num_objects: |
| padding = torch.zeros( |
| data.shape[0], |
| max_num_objects - data.shape[1], |
| 4, |
| dtype=data.dtype, |
| device=data.device, |
| ) |
| padding[:, :, 0] = 1.0 |
| data = torch.cat([data, padding], dim=1) |
| else: |
| data = data[:, :max_num_objects] |
|
|
| return data |
|
|
|
|
| def interpolate_contact_center( |
| contact_points_dict, |
| source_fps, |
| target_fps, |
| num_frames, |
| ): |
| """Compute contact center and in_contact label from raw contact points. |
| |
| Directly scales source frame indices to target frame space, avoiding |
| dense array interpolation that would blend real positions with zeros. |
| |
| Args: |
| contact_points_dict: Dict mapping frame_idx -> (N_points, 3) array |
| source_fps: Original frame rate |
| target_fps: Target frame rate |
| num_frames: Target number of frames |
| |
| Returns: |
| Tuple of: |
| contact_center: Tensor of shape (num_frames, 3) |
| in_contact: Tensor of shape (num_frames,) with binary labels |
| """ |
| if not contact_points_dict: |
| return torch.zeros(num_frames, 3), torch.zeros(num_frames) |
|
|
| fps_ratio = target_fps / source_fps |
| contact_center = torch.zeros(num_frames, 3) |
| in_contact = torch.zeros(num_frames) |
|
|
| for src_idx, points in contact_points_dict.items(): |
| if not (hasattr(points, "shape") and len(points) > 0): |
| continue |
| |
| t_start = max(0, int(src_idx * fps_ratio)) |
| t_end = min(num_frames, int((src_idx + 1) * fps_ratio) + 1) |
| center = torch.from_numpy(points.mean(axis=0).astype(np.float32)) |
| contact_center[t_start:t_end] = center |
| in_contact[t_start:t_end] = 1.0 |
|
|
| return contact_center, in_contact |
|
|
|
|
| class MotionLibBase: |
| def __init__(self, motion_lib_cfg, num_envs, device): |
| self.m_cfg = motion_lib_cfg |
| self.motion_fps_scale = self.m_cfg.get("motion_fps_scale", 1.0) |
| self._sim_fps = 1 / self.m_cfg.get("step_dt", 1 / 50) |
| self.target_fps = self.m_cfg.get("target_fps", 50) |
| self.adaptive_sampling_cfg = self.m_cfg.get("adaptive_sampling", {}) |
| self.all_motions_loaded = False |
|
|
| self.debug = motion_lib_cfg.get("debug", False) |
| self.use_parallel_fk = motion_lib_cfg.get("use_parallel_fk", False) |
| self.num_envs = num_envs |
| self._device = device |
| self.mesh_parsers = None |
| self.has_action = False |
| skeleton_file = Path(self.m_cfg.asset.assetRoot) / self.m_cfg.asset.assetFileName |
| self.skeleton_tree = skeleton.SkeletonTree.from_mjcf(skeleton_file) |
| logger.info(f"Loaded skeleton from {skeleton_file}") |
| logger.info(f"Loading motion data from {self.m_cfg.motion_file}...") |
| self.load_data(self.m_cfg.motion_file) |
| self.use_adaptive_sampling = self.adaptive_sampling_cfg.get("enable", False) |
| if self.use_adaptive_sampling: |
| self.init_adaptive_sampling() |
| self.setup_constants( |
| fix_height=motion_lib_cfg.get("fix_height", FixHeightMode.no_fix), |
| multi_thread=self.m_cfg.get("multi_thread", True), |
| ) |
|
|
| self.vid_smpl_pose = None |
| self.vid_smpl_joints = None |
| self.smpl_data = None |
| smpl_motion_file = motion_lib_cfg.get("smpl_motion_file", None) |
| self.smpl_data_keys = set() |
| if smpl_motion_file is not None: |
| if smpl_motion_file in ("dummy", "zeros"): |
| |
| |
| self.smpl_data = [None] * len(self._motion_data_keys) |
| elif osp.exists(smpl_motion_file): |
| if osp.isfile(smpl_motion_file): |
| self.smpl_data = joblib.load(smpl_motion_file) |
| self.smpl_data_keys = set(self.smpl_data.keys()) |
| self.smpl_data = [ |
| (self.smpl_data[k] if k in self.smpl_data else None) |
| for k in self._motion_data_keys |
| ] |
| else: |
| self.smpl_data = [] |
| smpl_pkl_files = set( |
| glob.glob(osp.join(smpl_motion_file, "**", "*.pkl"), recursive=True) |
| ) |
| for k in self._motion_data_keys: |
| seq = os.path.basename(k) |
| smpl_path = osp.join(smpl_motion_file, seq + ".pkl") |
| if self.debug or smpl_path in smpl_pkl_files: |
| self.smpl_data.append({"seq": seq, "path": smpl_path}) |
| self.smpl_data_keys.add(seq) |
| else: |
| self.smpl_data.append(None) |
| else: |
| self.smpl_data = [None] * len(self._motion_data_keys) |
|
|
| self.smpl_y_up = motion_lib_cfg.get("smpl_y_up", False) |
|
|
| |
| self.soma_data = None |
| soma_motion_file = motion_lib_cfg.get("soma_motion_file", None) |
| self.soma_data_keys = set() |
| self.soma_y_up = motion_lib_cfg.get("soma_y_up", True) |
| self.num_soma_joints = motion_lib_cfg.get("num_soma_joints", 26) |
| if soma_motion_file is not None: |
| if soma_motion_file in ("dummy", "zeros"): |
| self.soma_data = [None] * len(self._motion_data_keys) |
| elif osp.exists(soma_motion_file): |
| if osp.isfile(soma_motion_file): |
| self.soma_data = joblib.load(soma_motion_file) |
| self.soma_data_keys = set(self.soma_data.keys()) |
| self.soma_data = [ |
| (self.soma_data[k] if k in self.soma_data else None) |
| for k in self._motion_data_keys |
| ] |
| else: |
| |
| soma_index = { |
| osp.splitext(osp.basename(f))[0]: f |
| for f in glob.glob( |
| osp.join(soma_motion_file, "**", "*.pkl"), recursive=True |
| ) |
| } |
| self.soma_data = [] |
| for k in self._motion_data_keys: |
| seq = os.path.basename(k) |
| soma_path = soma_index.get(seq) |
| if soma_path is not None or self.debug: |
| self.soma_data.append( |
| { |
| "seq": seq, |
| "path": soma_path or osp.join(soma_motion_file, seq + ".pkl"), |
| } |
| ) |
| self.soma_data_keys.add(seq) |
| else: |
| self.soma_data.append(None) |
| else: |
| self.soma_data = [None] * len(self._motion_data_keys) |
|
|
| |
| self.object_data = None |
| object_motion_file = motion_lib_cfg.get("object_motion_file", None) |
| self.object_data_keys = set() |
| self.max_num_objects = motion_lib_cfg.get("max_num_objects", 1) |
| if object_motion_file is not None: |
| if osp.isfile(object_motion_file): |
| self.object_data = joblib.load(object_motion_file) |
| self.object_data_keys = set(self.object_data.keys()) |
| self.object_data = [ |
| (self.object_data[k] if k in self.object_data else None) |
| for k in self._motion_data_keys |
| ] |
| else: |
| self.object_data = [] |
| |
| |
| for k in self._motion_data_keys: |
| seq = os.path.basename(k) |
| object_path = osp.join(object_motion_file, seq + ".pkl") |
| if self.debug or osp.exists(object_path): |
| self.object_data.append({"seq": seq, "path": object_path}) |
| self.object_data_keys.add(seq) |
| else: |
| self.object_data.append(None) |
| |
| self.randomize_upper_body_poses = self.m_cfg.get("cat_upper_body_poses", False) |
| self.cat_upper_body_poses_prob = self.m_cfg.get("cat_upper_body_poses_prob", 0.0) |
| |
| self.upper_body_augment_prefixes = self.m_cfg.get( |
| "upper_body_augment_prefixes", |
| ["2025", "walking_2025", "running_2025", "slow_walk_2025"], |
| ) |
| |
| self.randomize_wrist_poses = self.m_cfg.get("randomize_wrist_poses", False) |
| self.randomize_wrist_prob = self.m_cfg.get("randomize_wrist_prob", 0.3) |
| self.randomize_wrist_std = self.m_cfg.get("randomize_wrist_std", 0.1) |
| |
| self.wrist_mujoco_dof_indices = [19, 20, 21, 26, 27, 28] |
|
|
| def load_data(self, motion_file): |
| if osp.isfile(motion_file): |
| self.mode = MotionlibMode.file |
| self._motion_data_load = joblib.load(motion_file) |
| else: |
| assert osp.isdir( |
| motion_file |
| ), f"Expected motion_file to be a directory, got: {motion_file}" |
| self.mode = MotionlibMode.directory |
| if self.debug: |
| self._motion_data_load = {} |
| else: |
| self._motion_data_load = { |
| osp.splitext(osp.basename(f))[0]: {"path": f} |
| for f in glob.glob(osp.join(motion_file, "**", "*.pkl"), recursive=True) |
| if not f.endswith("metadata.pkl") |
| } |
|
|
| metadata_files = [] |
| |
| direct_metadata = osp.join(motion_file, "metadata.pkl") |
| if osp.exists(direct_metadata): |
| metadata_files.append(direct_metadata) |
| |
| all_sub_dirs = os.listdir(motion_file) |
| if self.debug: |
| all_sub_dirs = all_sub_dirs[:1] |
| for sub_dir in all_sub_dirs: |
| sub_dir_path = osp.join(motion_file, sub_dir) |
| if osp.isdir(sub_dir_path): |
| sub_meta = osp.join(sub_dir_path, "metadata.pkl") |
| if osp.exists(sub_meta): |
| metadata_files.append(sub_meta) |
|
|
| for metadata_file in metadata_files: |
| metadata = joblib.load(metadata_file) |
| if self.debug: |
| metadata = { |
| k: v |
| for k, v in list(metadata.items())[:5] |
| if osp.exists(f"{sub_dir_path}/{k}.pkl") |
| } |
| for k, v in metadata.items(): |
| if self.debug: |
| self._motion_data_load[k] = {"path": f"{sub_dir_path}/{k}.pkl"} |
| if ( |
| k in self._motion_data_load |
| ): |
| self._motion_data_load[k].update(v) |
|
|
| print(f"Loaded {len(self._motion_data_load)} motion files") |
|
|
| data_list = self._motion_data_load |
|
|
| filter_motion_keys = self.m_cfg.get("filter_motion_keys", None) |
| if filter_motion_keys is not None: |
| if isinstance(filter_motion_keys, str): |
| patterns = [filter_motion_keys] |
| else: |
| patterns = list(filter_motion_keys) |
|
|
| if all(pattern in data_list for pattern in patterns): |
| matched_keys = [pattern for pattern in patterns if pattern in data_list] |
| else: |
| compiled = [] |
| for pattern in patterns: |
| try: |
| compiled.append(re.compile(pattern)) |
| except re.error as exc: |
| raise ValueError(f"Invalid filter_motion_keys regex: {pattern}") from exc |
|
|
| matched_keys = [ |
| k for k in data_list if any(regex.fullmatch(k) for regex in compiled) |
| ] |
| matched_keys.sort() |
| data_list = {k: data_list[k] for k in matched_keys} |
|
|
| remove_motion_keys = self.m_cfg.get("remove_motion_keys", None) |
| if remove_motion_keys is not None: |
| |
| keys_to_remove = [ |
| k for k in data_list if any(k.startswith(prefix) for prefix in remove_motion_keys) |
| ] |
| for k in keys_to_remove: |
| del data_list[k] |
|
|
| max_unique_motions = self.m_cfg.get("max_unique_motions", None) |
| if max_unique_motions is not None and len(data_list) > max_unique_motions: |
| import random |
|
|
| keys = sorted(data_list.keys()) |
| selected = random.sample(keys, max_unique_motions) |
| data_list = {k: data_list[k] for k in selected} |
| print( |
| f"Limited to {max_unique_motions} random motions (from {len(keys)})" |
| ) |
|
|
| self._motion_data_list = np.array(list(data_list.values())) |
| self._motion_data_keys = np.array(list(data_list.keys())) |
|
|
| |
| |
| |
| |
| |
| |
|
|
| self._num_unique_motions = len(self._motion_data_list) |
| logger.info(f"Loaded {self._num_unique_motions} motions") |
|
|
| def _should_augment_upper_body(self, motion_key): |
| """Check if motion key matches any prefix for upper body augmentation""" |
| return any(motion_key.startswith(prefix) for prefix in self.upper_body_augment_prefixes) |
|
|
| def setup_constants(self, fix_height=FixHeightMode.full_fix, multi_thread=True): |
| self.fix_height = fix_height |
| self.multi_thread = multi_thread |
|
|
| |
| self._curr_motion_ids = None |
| self._termination_history = torch.zeros(self._num_unique_motions).to(self._device) |
| self._success_rate = torch.zeros(self._num_unique_motions).to(self._device) |
| self._sampling_history = torch.zeros(self._num_unique_motions).to(self._device) |
| self._sampling_prob = ( |
| torch.ones(self._num_unique_motions).to(self._device) / self._num_unique_motions |
| ) |
|
|
| def update_soft_sampling_weight(self, failed_keys): |
| |
| if len(failed_keys) > 0: |
| all_keys = self._motion_data_keys.tolist() |
| indexes = [all_keys.index(k) for k in failed_keys] |
| self._termination_history[indexes] += 1 |
| self.update_sampling_prob(self._termination_history) |
|
|
| print( |
| "############################################################ Auto PMCP ############################################################" |
| ) |
| print( |
| f"Training mostly on {len(self._sampling_prob.cpu().nonzero())} seqs " |
| ) |
| print( |
| self._motion_data_keys[self._sampling_prob.cpu().nonzero()].flatten() |
| ) |
| print( |
| "###############################################################################################################################" |
| ) |
| else: |
| all_keys = self._motion_data_keys.tolist() |
| self._sampling_prob = ( |
| torch.ones(self._num_unique_motions).to(self._device) / self._num_unique_motions |
| ) |
|
|
| def update_sampling_prob(self, termination_history): |
| if ( |
| len(termination_history) == len(self._termination_history) |
| and termination_history.sum() > 0 |
| ): |
| self._sampling_prob[:] = termination_history / termination_history.sum() |
| if self._sampling_prob[self._curr_motion_ids].sum() == 0: |
| self._sampling_prob[self._curr_motion_ids] += 1e-6 |
| self._sampling_prob[:] = self._sampling_prob[:] / self._sampling_prob[:].sum() |
| self._sampling_batch_prob = ( |
| self._sampling_prob[self._curr_motion_ids] |
| / self._sampling_prob[self._curr_motion_ids].sum() |
| ) |
| self._termination_history = termination_history |
| return True |
| else: |
| return False |
|
|
| def get_motion_actions(self, motion_ids, motion_times): |
| motion_len = self._motion_lengths[motion_ids] |
| num_frames = self._motion_num_frames[motion_ids] |
| dt = self._motion_dt[motion_ids] |
| |
| frame_idx0, frame_idx1, blend = self._calc_frame_blend( |
| motion_times, motion_len, num_frames, dt |
| ) |
| f0l = frame_idx0 + self.length_starts[motion_ids] |
| f1l = frame_idx1 + self.length_starts[motion_ids] |
|
|
| action = self._motion_actions[f0l] |
| return action |
|
|
| def get_time_step_total(self, motion_ids): |
| return self._motion_num_frames[motion_ids] |
|
|
| @property |
| def body_indexes(self): |
| return self.m_cfg.get("body_indexes_data", None) |
|
|
| def get_dof_pos(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
|
|
| return self.dof_pos[motion_steps + length_starts] |
|
|
| def get_dof_vel(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self.dof_vel[motion_steps + length_starts] |
|
|
| def get_hand_dof_pos(self, motion_ids, motion_steps): |
| """Get hand DOF positions if available (for 43-DOF motion on 43-DOF robot).""" |
| if not hasattr(self, "hand_dof_pos") or self.hand_dof_pos is None: |
| return None |
| length_starts = self.length_starts[motion_ids] |
| return self.hand_dof_pos[motion_steps + length_starts] |
|
|
| def get_body_pos_w(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
|
|
| return self.body_pos_w[motion_steps + length_starts] |
|
|
| def get_body_quat_w(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self.body_quat_w[motion_steps + length_starts] |
|
|
| def get_body_lin_vel_w(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self.body_lin_vel_w[motion_steps + length_starts] |
|
|
| def get_body_ang_vel_w(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self.body_ang_vel_w[motion_steps + length_starts] |
|
|
| |
| def get_body_pos_w_full(self, motion_ids, motion_steps): |
| """Get full body positions (all bodies, IsaacLab order).""" |
| length_starts = self.length_starts[motion_ids] |
| return self.body_pos_w_full[motion_steps + length_starts] |
|
|
| def get_body_quat_w_full(self, motion_ids, motion_steps): |
| """Get full body quaternions (all bodies, IsaacLab order, wxyz).""" |
| length_starts = self.length_starts[motion_ids] |
| return self.body_quat_w_full[motion_steps + length_starts] |
|
|
| def get_body_lin_vel_w_full(self, motion_ids, motion_steps): |
| """Get full body linear velocities (all bodies, IsaacLab order).""" |
| length_starts = self.length_starts[motion_ids] |
| return self.body_lin_vel_w_full[motion_steps + length_starts] |
|
|
| def get_body_ang_vel_w_full(self, motion_ids, motion_steps): |
| """Get full body angular velocities (all bodies, IsaacLab order).""" |
| length_starts = self.length_starts[motion_ids] |
| return self.body_ang_vel_w_full[motion_steps + length_starts] |
|
|
| def get_root_pos_w(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self.body_pos_w[motion_steps + length_starts, 0, :] |
|
|
| def get_root_quat_w(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self.body_quat_w[motion_steps + length_starts, 0, :] |
|
|
| def get_root_lin_vel_w(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self.body_lin_vel_w[motion_steps + length_starts, 0, :] |
|
|
| def get_root_ang_vel_w(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self.body_ang_vel_w[motion_steps + length_starts, 0, :] |
|
|
| def get_smpl_pose(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self._motion_smpl_poses[motion_steps + length_starts] |
|
|
| def get_smpl_joints(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self._motion_smpl_joints[motion_steps + length_starts] |
|
|
| def get_smpl_transl(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self._motion_smpl_transl[motion_steps + length_starts] |
|
|
| @staticmethod |
| def _resample_soma_tensor(data, fps_source, fps_target): |
| """Resample a SOMA tensor along dim 0 matching interploate_pose frame count. |
| |
| Uses the same arange(0, duration, 1/fps_target) formula as |
| torch_humanoid_batch.interploate_pose so robot and SOMA frame counts align. |
| """ |
| n_src = data.shape[0] |
| duration = (n_src - 1) / fps_source |
| tgt_times = torch.arange(0, duration, 1 / fps_target, dtype=torch.float32) |
| n_tgt = len(tgt_times) |
| if n_tgt <= 1: |
| return data[:1] |
| |
| phase = tgt_times / duration |
| idx0 = (phase * (n_src - 1)).floor().long() |
| idx1 = torch.minimum(idx0 + 1, torch.tensor(n_src - 1)) |
| blend = (phase * (n_src - 1) - idx0).float() |
| |
| for _ in range(data.dim() - 1): |
| blend = blend.unsqueeze(-1) |
| return data[idx0] * (1 - blend) + data[idx1] * blend |
|
|
| def get_soma_joints(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self._motion_soma_joints[motion_steps + length_starts] |
|
|
| def get_soma_root_quat(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self._motion_soma_root_quat[motion_steps + length_starts] |
|
|
| def get_soma_transl(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self._motion_soma_transl[motion_steps + length_starts] |
|
|
| def get_object_root_pos(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self._motion_object_root_pos[motion_steps + length_starts] |
|
|
| def get_object_root_quat(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self._motion_object_root_quat[motion_steps + length_starts] |
|
|
| def get_object_lin_vel(self, motion_ids, motion_steps): |
| """Get object linear velocity from motion library.""" |
| length_starts = self.length_starts[motion_ids] |
| return self._motion_object_lin_vel[motion_steps + length_starts] |
|
|
| def get_object_ang_vel(self, motion_ids, motion_steps): |
| """Get object angular velocity from motion library.""" |
| length_starts = self.length_starts[motion_ids] |
| return self._motion_object_ang_vel[motion_steps + length_starts] |
|
|
| def get_object_contact_center(self, motion_ids, motion_steps, hand="right_hand"): |
| """Get object contact center from motion library. |
| |
| Contact center is the mean of all contact points per frame for the given hand. |
| |
| Args: |
| motion_ids: (N,) tensor of motion indices |
| motion_steps: (N,) tensor of frame indices within each motion |
| hand: Which hand's contact center to return ("left_hand" or "right_hand") |
| |
| Returns: |
| Tensor of shape (N, 3) with contact center positions in object-local frame, |
| or None if not available. |
| """ |
| attr = f"_motion_object_contact_center_{'left' if hand == 'left_hand' else 'right'}" |
| if not hasattr(self, attr): |
| return None |
| length_starts = self.length_starts[motion_ids] |
| return getattr(self, attr)[motion_steps + length_starts] |
|
|
| def get_object_in_contact(self, motion_ids, motion_steps, hand="right_hand"): |
| """Get binary in_contact label for the given hand. |
| |
| Args: |
| motion_ids: (N,) tensor of motion indices |
| motion_steps: (N,) tensor of frame indices within each motion |
| hand: Which hand ("left_hand" or "right_hand") |
| |
| Returns: |
| Tensor of shape (N,) with 1.0 if in contact, 0.0 otherwise, |
| or None if not available. |
| """ |
| attr = f"_motion_object_in_contact_{'left' if hand == 'left_hand' else 'right'}" |
| if not hasattr(self, attr): |
| return None |
| length_starts = self.length_starts[motion_ids] |
| return getattr(self, attr)[motion_steps + length_starts] |
|
|
| def get_hand_action(self, motion_ids, motion_steps, hand="right_hand"): |
| """Get discrete hand action (open/closed) for the given hand. |
| |
| Args: |
| motion_ids: (N,) tensor of motion indices |
| motion_steps: (N,) tensor of frame indices within each motion |
| hand: Which hand ("left_hand" or "right_hand") |
| |
| Returns: |
| Tensor of shape (N,) with -1.0 = open, +1.0 = closed, |
| or None if not available. |
| """ |
| attr = f"_motion_hand_action_{'left' if hand == 'left_hand' else 'right'}" |
| if not hasattr(self, attr): |
| return None |
| length_starts = self.length_starts[motion_ids] |
| return getattr(self, attr)[motion_steps + length_starts] |
|
|
| def get_feet_l(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self.feet_l[motion_steps + length_starts] |
|
|
| def get_feet_r(self, motion_ids, motion_steps): |
| length_starts = self.length_starts[motion_ids] |
| return self.feet_r[motion_steps + length_starts] |
|
|
| def get_motion_state(self, motion_ids, motion_times, offset=None): |
| motion_len = self._motion_lengths[motion_ids] |
| num_frames = self._motion_num_frames[motion_ids] |
| dt = self._motion_dt[motion_ids] |
|
|
| frame_idx0, frame_idx1, blend = self._calc_frame_blend( |
| motion_times, motion_len, num_frames, dt |
| ) |
| f0l = frame_idx0 + self.length_starts[motion_ids] |
| f1l = frame_idx1 + self.length_starts[motion_ids] |
|
|
| if "dof_pos" in self.__dict__: |
| local_rot0 = self.dof_pos[f0l] |
| local_rot1 = self.dof_pos[f1l] |
| else: |
| local_rot0 = self.body_pos_b[f0l] |
| local_rot1 = self.body_pos_b[f1l] |
|
|
| body_lin_vel_w0 = self.body_lin_vel_w[f0l] |
| body_lin_vel_w1 = self.body_lin_vel_w[f1l] |
|
|
| body_ang_vel0 = self.body_ang_vel_w[f0l] |
| body_ang_vel1 = self.body_ang_vel_w[f1l] |
|
|
| body_pos_w0 = self.body_pos_w[f0l, :] |
| body_pos_w1 = self.body_pos_w[f1l, :] |
|
|
| dof_vel0 = self.dof_vel[f0l] |
| dof_vel1 = self.dof_vel[f1l] |
|
|
| vals = [ |
| local_rot0, |
| local_rot1, |
| body_lin_vel_w0, |
| body_lin_vel_w1, |
| body_ang_vel0, |
| body_ang_vel1, |
| body_pos_w0, |
| body_pos_w1, |
| dof_vel0, |
| dof_vel1, |
| ] |
| for v in vals: |
| assert v.dtype != torch.float64 |
|
|
| blend = blend.unsqueeze(-1) |
|
|
| blend_exp = blend.unsqueeze(-1) |
|
|
| if offset is None: |
| body_pos_w = ( |
| 1.0 - blend_exp |
| ) * body_pos_w0 + blend_exp * body_pos_w1 |
| else: |
| body_pos_w = ( |
| (1.0 - blend_exp) * body_pos_w0 + blend_exp * body_pos_w1 + offset[..., None, :] |
| ) |
|
|
| body_lin_vel_w = (1.0 - blend_exp) * body_lin_vel_w0 + blend_exp * body_lin_vel_w1 |
| body_ang_vel_w = (1.0 - blend_exp) * body_ang_vel0 + blend_exp * body_ang_vel1 |
|
|
| if "dof_pos" in self.__dict__: |
| dof_vel = (1.0 - blend) * dof_vel0 + blend * dof_vel1 |
| dof_pos = (1.0 - blend) * local_rot0 + blend * local_rot1 |
| else: |
| dof_vel = (1.0 - blend_exp) * dof_vel0 + blend_exp * dof_vel1 |
| local_rot = rotations.slerp(local_rot0, local_rot1, torch.unsqueeze(blend, axis=-1)) |
| dof_pos = self._local_rotation_to_dof_smpl(local_rot) |
|
|
| body_quat_w0 = self.body_quat_w[f0l] |
| body_quat_w1 = self.body_quat_w[f1l] |
| body_quat_w = rotations.slerp(body_quat_w0, body_quat_w1, blend_exp) |
| return_dict = {} |
|
|
| if "gts_t" in self.__dict__: |
| body_pos_w_t0 = self.body_pos_t_w[f0l] |
| body_pos_w_t1 = self.body_pos_t_w[f1l] |
|
|
| body_quat_t0 = self.body_quat_t_w[f0l] |
| body_quat_t1 = self.body_quat_t_w[f1l] |
|
|
| body_lin_vel_w_t0 = self.body_lin_vel_t_w[f0l] |
| body_lin_vel_w_t1 = self.body_lin_vel_t_w[f1l] |
|
|
| body_ang_vel_t0 = self.body_ang_vel_t_w[f0l] |
| body_ang_vel_t1 = self.body_ang_vel_t_w[f1l] |
| if offset is None: |
| body_pos_t_w = (1.0 - blend_exp) * body_pos_w_t0 + blend_exp * body_pos_w_t1 |
| else: |
| body_pos_t_w = ( |
| (1.0 - blend_exp) * body_pos_w_t0 |
| + blend_exp * body_pos_w_t1 |
| + offset[..., None, :] |
| ) |
| body_quat_t_w = rotations.slerp(body_quat_t0, body_quat_t1, blend_exp) |
| body_lin_vel_t_w = (1.0 - blend_exp) * body_lin_vel_w_t0 + blend_exp * body_lin_vel_w_t1 |
| body_ang_vel_t_w = (1.0 - blend_exp) * body_ang_vel_t0 + blend_exp * body_ang_vel_t1 |
| else: |
| body_pos_t_w = body_pos_w |
| body_quat_t_w = body_quat_w |
| body_lin_vel_t_w = body_lin_vel_w |
| body_ang_vel_t_w = body_ang_vel_w |
|
|
| if self.smpl_data is not None: |
| smpl_pose0 = self._motion_smpl_poses[f0l] |
| smpl_pose1 = self._motion_smpl_poses[f1l] |
| smpl_pose = (1.0 - blend) * smpl_pose0 + blend * smpl_pose1 |
| return_dict.update({"smpl_pose": smpl_pose.clone()}) |
|
|
| if hasattr(self, "_motion_smpl_joints"): |
| smpl_joints0 = self._motion_smpl_joints[f0l] |
| smpl_joints1 = self._motion_smpl_joints[f1l] |
| smpl_joints = (1.0 - blend_exp) * smpl_joints0 + blend_exp * smpl_joints1 |
| return_dict.update({"smpl_joints": smpl_joints.clone()}) |
|
|
| if hasattr(self, "_motion_smpl_transl"): |
| smpl_transl0 = self._motion_smpl_transl[f0l] |
| smpl_transl1 = self._motion_smpl_transl[f1l] |
| smpl_transl = (1.0 - blend_exp) * smpl_transl0 + blend_exp * smpl_transl1 |
| return_dict.update({"smpl_transl": smpl_transl.clone()}) |
|
|
| if self.soma_data is not None: |
| if hasattr(self, "_motion_soma_joints"): |
| soma_joints0 = self._motion_soma_joints[f0l] |
| soma_joints1 = self._motion_soma_joints[f1l] |
| soma_joints = (1.0 - blend_exp) * soma_joints0 + blend_exp * soma_joints1 |
| return_dict.update({"soma_joints": soma_joints.clone()}) |
|
|
| if hasattr(self, "_motion_soma_root_quat"): |
| |
| soma_rq0 = self._motion_soma_root_quat[f0l] |
| soma_rq1 = self._motion_soma_root_quat[f1l] |
| soma_root_quat = (1.0 - blend_exp) * soma_rq0 + blend_exp * soma_rq1 |
| soma_root_quat = soma_root_quat / (soma_root_quat.norm(dim=-1, keepdim=True) + 1e-8) |
| return_dict.update({"soma_root_quat": soma_root_quat.clone()}) |
|
|
| if hasattr(self, "_motion_soma_transl"): |
| soma_transl0 = self._motion_soma_transl[f0l] |
| soma_transl1 = self._motion_soma_transl[f1l] |
| soma_transl = (1.0 - blend_exp) * soma_transl0 + blend_exp * soma_transl1 |
| return_dict.update({"soma_transl": soma_transl.clone()}) |
|
|
| if self.object_data is not None: |
| if hasattr(self, "_motion_object_root_pos"): |
| object_root_pos0 = self._motion_object_root_pos[f0l] |
| object_root_pos1 = self._motion_object_root_pos[f1l] |
| object_root_pos = ( |
| 1.0 - blend_exp |
| ) * object_root_pos0 + blend_exp * object_root_pos1 |
| return_dict.update({"object_root_pos": object_root_pos.clone()}) |
|
|
| if hasattr(self, "_motion_object_root_quat"): |
| object_root_quat0 = self._motion_object_root_quat[f0l] |
| object_root_quat1 = self._motion_object_root_quat[f1l] |
| |
| object_root_quat = rotations.slerp(object_root_quat0, object_root_quat1, blend_exp) |
| return_dict.update({"object_root_quat": object_root_quat.clone()}) |
|
|
| return_dict.update( |
| { |
| "root_pos": body_pos_w[..., 0, :].clone(), |
| "root_rot": body_quat_w[..., 0, :].clone(), |
| "dof_pos": dof_pos.clone(), |
| "root_vel": body_lin_vel_w[..., 0, :].clone(), |
| "root_ang_vel": body_ang_vel_w[..., 0, :].clone(), |
| "dof_vel": dof_vel.clone(), |
| "motion_aa": self._motion_aa[f0l].clone(), |
| "motion_bodies": self._motion_bodies[motion_ids].clone(), |
| "body_pos_w": body_pos_w.clone(), |
| "body_quat_w": body_quat_w.clone(), |
| "body_lin_vel_w": body_lin_vel_w.clone(), |
| "body_ang_vel_w": body_ang_vel_w.clone(), |
| "body_pos_w_t": body_pos_t_w.clone(), |
| "body_quat_t": body_quat_t_w.clone(), |
| "body_lin_vel_w_t": body_lin_vel_t_w.clone(), |
| "body_ang_vel_t": body_ang_vel_t_w.clone(), |
| } |
| ) |
| if "feet_l" in self.__dict__: |
| blend_int = blend.round().int() |
| feet_l = torch.where(blend_int == 0, self.feet_l[f0l], self.feet_l[f1l]) |
| feet_r = torch.where(blend_int == 0, self.feet_r[f0l], self.feet_r[f1l]) |
| return_dict.update( |
| { |
| "feet_l": feet_l.clone().bool(), |
| "feet_r": feet_r.clone().bool(), |
| } |
| ) |
| return return_dict |
|
|
| def load_all_motions(self): |
| self.all_motions_loaded = True |
| self.load_motions(random_sample=False, num_motions_to_load=self._num_unique_motions) |
|
|
| def load_motions_for_training(self, max_num_seqs=None): |
| if self.all_motions_loaded: |
| print("All motions already loaded!!! No need to resample.") |
| return False |
|
|
| if self.m_cfg.get("override_num_motions_to_load", None) is not None: |
| max_num_seqs = self.m_cfg.override_num_motions_to_load |
|
|
| |
| load_unique = self.m_cfg.get("load_unique_motions", False) |
|
|
| if ( |
| max_num_seqs is None |
| ): |
| max_num_seqs = self._num_unique_motions |
| self.all_motions_loaded = True |
| self.load_motions(random_sample=False, num_motions_to_load=self._num_unique_motions) |
| elif ( |
| max_num_seqs >= self._num_unique_motions |
| ): |
| self.all_motions_loaded = True |
| self.load_motions(random_sample=False, num_motions_to_load=self._num_unique_motions) |
| else: |
| self.all_motions_loaded = False |
| |
| self.load_motions(random_sample=not load_unique, num_motions_to_load=max_num_seqs) |
| if load_unique: |
| print( |
| f"[MotionLib] Loaded {max_num_seqs} unique motions (no duplicates)" |
| ) |
| return True |
|
|
| def load_motions_for_evaluation(self, start_idx=0): |
| |
| |
| |
| |
|
|
| if ( |
| self._num_unique_motions > self.num_envs |
| ): |
| self.all_motions_loaded = False |
| self.load_motions( |
| random_sample=False, |
| num_motions_to_load=self.num_envs, |
| start_idx=start_idx, |
| is_evaluation=True, |
| ) |
| else: |
| self.all_motions_loaded = True |
| self.load_motions( |
| random_sample=False, |
| num_motions_to_load=self._num_unique_motions, |
| start_idx=start_idx, |
| is_evaluation=True, |
| ) |
|
|
| def load_motions( |
| self, |
| random_sample=True, |
| start_idx=0, |
| max_len=-1, |
| target_heading=None, |
| num_motions_to_load=None, |
| is_evaluation=False, |
| ): |
|
|
| if "gts" in self.__dict__: |
| del ( |
| self.body_pos_w, |
| self.body_quat_w, |
| self.body_pos_b, |
| self.root_linv_vel_w, |
| self.root_ang_vel_w, |
| self.body_ang_vel_w, |
| self.body_lin_vel_w, |
| self.dof_vels, |
| self.dof_pos, |
| ) |
| if "gts_t" in self.__dict__: |
| del ( |
| self.body_pos_t_w, |
| self.body_quat_t_w, |
| self.body_lin_vel_t_w, |
| self.body_ang_vel_t_w, |
| ) |
|
|
| motions = [] |
| _motion_lengths = [] |
| _motion_fps = [] |
| _motion_dt = [] |
| _motion_num_frames = [] |
| _motion_bodies = [] |
| _motion_aa = [] |
| has_action = False |
| _motion_actions = [] |
| _motion_smpl_poses = [] |
| _motion_smpl_joints = [] |
| _motion_smpl_transl = [] |
| _motion_soma_joints = [] |
| _motion_soma_root_quat = [] |
| _motion_soma_transl = [] |
| _motion_object_root_pos = [] |
| _motion_object_root_quat = [] |
| _motion_object_contact_center_left = [] |
| _motion_object_contact_center_right = [] |
| _motion_object_in_contact_left = [] |
| _motion_object_in_contact_right = [] |
| _motion_hand_action_left = [] |
| _motion_hand_action_right = [] |
|
|
| total_len = 0.0 |
| self.num_joints = len(self.skeleton_tree.node_names) |
| if num_motions_to_load is None: |
| num_motion_to_load = self.num_envs |
| else: |
| num_motion_to_load = num_motions_to_load |
|
|
| if self.use_adaptive_sampling: |
| self.update_adaptive_sampling_motion_sequences() |
|
|
| if random_sample: |
| sample_idxes = torch.multinomial( |
| self._sampling_prob, num_samples=num_motion_to_load, replacement=True |
| ).to(self._device) |
| else: |
| sample_idxes = torch.clamp( |
| torch.arange(num_motion_to_load) + start_idx, max=self._num_unique_motions - 1 |
| ).to(self._device) |
|
|
| |
| self._curr_motion_ids = sample_idxes |
| self.curr_motion_keys = ( |
| [self._motion_data_keys[sample_idxes.cpu()]] |
| if sample_idxes.numel() == 1 |
| else self._motion_data_keys[sample_idxes.cpu()].tolist() |
| ) |
| self._sampling_batch_prob = ( |
| self._sampling_prob[self._curr_motion_ids] |
| / self._sampling_prob[self._curr_motion_ids].sum() |
| ) |
|
|
| logger.info(f"Loading {num_motion_to_load} motions...") |
| logger.info(f"Sampling motion: {sample_idxes[:10]}, ....") |
| logger.info(f"Current motion keys: {self.curr_motion_keys[:10]}, ....") |
|
|
| motion_data_list = self._motion_data_list[sample_idxes.cpu().numpy()] |
| if self.smpl_data is not None: |
| smpl_data_list = [self.smpl_data[idx] for idx in sample_idxes.cpu().numpy()] |
| else: |
| smpl_data_list = None |
| if self.object_data is not None: |
| object_data_list = [self.object_data[idx] for idx in sample_idxes.cpu().numpy()] |
| else: |
| object_data_list = None |
| if self.soma_data is not None: |
| soma_data_list = [self.soma_data[idx] for idx in sample_idxes.cpu().numpy()] |
| else: |
| soma_data_list = None |
| torch.set_num_threads(1) |
|
|
| |
| try: |
| soft_limit, hard_limit = resource.getrlimit(resource.RLIMIT_NOFILE) |
| target_limit = 1048576 |
|
|
| |
| if soft_limit < target_limit: |
| try: |
| |
| resource.setrlimit(resource.RLIMIT_NOFILE, (target_limit, target_limit)) |
| logger.info( |
| f"Increased file descriptor limits from {soft_limit}/{hard_limit} to {target_limit}/{target_limit}" |
| ) |
| except PermissionError: |
| |
| new_soft = min(target_limit, hard_limit) |
| resource.setrlimit(resource.RLIMIT_NOFILE, (new_soft, hard_limit)) |
| logger.info( |
| f"Increased soft file descriptor limit from {soft_limit} to {new_soft} (hard limit: {hard_limit})" |
| ) |
| except Exception as e: |
| logger.warning(f"Could not increase file descriptor limit: {e}") |
|
|
| manager = mp.Manager() |
| queue = manager.Queue() |
| num_jobs = min(min(mp.cpu_count(), 32), len(motion_data_list)) |
|
|
| if num_jobs <= 8 or not self.multi_thread or len(motion_data_list) <= 128: |
| num_jobs = 1 |
|
|
| logger.info(f"Loading motions with {num_jobs} jobs...") |
| self.res_non_nav_dataset = {} |
| res_acc = {} |
| workers = [] |
|
|
| |
| |
| if self.randomize_upper_body_poses and not is_evaluation: |
|
|
| |
| nav_indices = [ |
| i |
| for i in range(len(motion_data_list)) |
| if self._should_augment_upper_body(self.curr_motion_keys[i]) |
| ] |
| other_indices = [ |
| i |
| for i in range(len(motion_data_list)) |
| if not self._should_augment_upper_body(self.curr_motion_keys[i]) |
| ] |
|
|
| nav_motion_data_list = [motion_data_list[i] for i in nav_indices] |
| other_motion_data_list = [motion_data_list[i] for i in other_indices] |
|
|
| if self.smpl_data is not None: |
| nav_smpl_data_list = [smpl_data_list[i] for i in nav_indices] |
| other_smpl_data_list = [smpl_data_list[i] for i in other_indices] |
| else: |
| nav_smpl_data_list = None |
| other_smpl_data_list = None |
|
|
| if self.object_data is not None: |
| nav_object_data_list = [object_data_list[i] for i in nav_indices] |
| other_object_data_list = [object_data_list[i] for i in other_indices] |
| else: |
| nav_object_data_list = None |
| other_object_data_list = None |
|
|
| if soma_data_list is not None: |
| nav_soma_data_list = [soma_data_list[i] for i in nav_indices] |
| other_soma_data_list = [soma_data_list[i] for i in other_indices] |
| else: |
| nav_soma_data_list = None |
| other_soma_data_list = None |
|
|
| |
| if len(other_motion_data_list) > 0: |
| jobs = other_motion_data_list |
| chunk = np.ceil(len(jobs) / num_jobs).astype(int) |
| ids = np.array(other_indices) |
|
|
| jobs = [ |
| ( |
| ids[i : i + chunk], |
| jobs[i : i + chunk], |
| ( |
| None |
| if other_smpl_data_list is None |
| else other_smpl_data_list[i : i + chunk] |
| ), |
| ( |
| None |
| if other_object_data_list is None |
| else other_object_data_list[i : i + chunk] |
| ), |
| ( |
| None |
| if other_soma_data_list is None |
| else other_soma_data_list[i : i + chunk] |
| ), |
| self.fix_height, |
| target_heading, |
| max_len, |
| is_evaluation, |
| ) |
| for i in range(0, len(jobs), chunk) |
| ] |
|
|
| job_args = [jobs[i] for i in range(len(jobs))] |
| for i in range(1, len(jobs)): |
| worker_args = (*job_args[i], queue, i) |
| worker = mp.Process(target=self.load_motion_with_skeleton, args=worker_args) |
| worker.start() |
| workers.append(worker) |
| res_acc.update(self.load_motion_with_skeleton(*jobs[0], None, 0)) |
|
|
| |
| for worker in workers: |
| worker.join() |
| worker.close() |
| workers = [] |
|
|
| for i in progress.track( |
| range(len(jobs) - 1), "Gathering results for non-navigation dataset..." |
| ): |
| res = queue.get() |
| res_acc.update(res) |
|
|
| self.res_non_nav_dataset = res_acc.copy() |
|
|
| |
| if len(nav_motion_data_list) > 0: |
|
|
| jobs = nav_motion_data_list |
| chunk = np.ceil(len(jobs) / num_jobs).astype(int) |
| ids = np.array(nav_indices) |
|
|
| jobs = [ |
| ( |
| ids[i : i + chunk], |
| jobs[i : i + chunk], |
| nav_smpl_data_list[ |
| i : i + chunk |
| ], |
| ( |
| None |
| if nav_object_data_list is None |
| else nav_object_data_list[i : i + chunk] |
| ), |
| (None if nav_soma_data_list is None else nav_soma_data_list[i : i + chunk]), |
| self.fix_height, |
| target_heading, |
| max_len, |
| is_evaluation, |
| ) |
| for i in range(0, len(jobs), chunk) |
| ] |
| job_args = [jobs[i] for i in range(len(jobs))] |
| for i in range(1, len(jobs)): |
| worker_args = (*job_args[i], queue, i) |
| worker = mp.Process(target=self.load_motion_with_skeleton, args=worker_args) |
| worker.start() |
| workers.append(worker) |
| res_acc.update(self.load_motion_with_skeleton(*jobs[0], None, 0)) |
|
|
| for i in progress.track( |
| range(len(jobs) - 1), "Gathering results for navigation dataset..." |
| ): |
| res = queue.get() |
| res_acc.update(res) |
|
|
| |
| for worker in workers: |
| worker.join() |
| worker.close() |
| workers = [] |
|
|
| else: |
| jobs = motion_data_list |
| chunk = np.ceil(len(jobs) / num_jobs).astype(int) |
| ids = np.arange(len(jobs)) |
|
|
| jobs = [ |
| ( |
| ids[i : i + chunk], |
| jobs[i : i + chunk], |
| None if smpl_data_list is None else smpl_data_list[i : i + chunk], |
| None if object_data_list is None else object_data_list[i : i + chunk], |
| None if soma_data_list is None else soma_data_list[i : i + chunk], |
| self.fix_height, |
| target_heading, |
| max_len, |
| is_evaluation, |
| ) |
| for i in range(0, len(jobs), chunk) |
| ] |
|
|
| job_args = [jobs[i] for i in range(len(jobs))] |
| for i in range(1, len(jobs)): |
| worker_args = (*job_args[i], queue, i) |
| worker = mp.Process(target=self.load_motion_with_skeleton, args=worker_args) |
| worker.start() |
| workers.append(worker) |
| res_acc.update(self.load_motion_with_skeleton(*jobs[0], None, 0)) |
|
|
| for i in progress.track(range(len(jobs) - 1), "Gathering results..."): |
| res = queue.get() |
| res_acc.update(res) |
|
|
| nav_indices = [] |
| other_indices = list(range(len(motions))) |
|
|
| |
| for worker in workers: |
| worker.join() |
| worker.close() |
| workers = [] |
|
|
| for f in progress.track(range(len(res_acc)), description="Processing motions..."): |
| motion_file_data, curr_motion = res_acc[f] |
| motion_fps = int(curr_motion.fps * self.motion_fps_scale) |
| curr_dt = 1.0 / motion_fps |
| num_frames = curr_motion.global_rotation.shape[0] |
|
|
| curr_len = 1.0 / motion_fps * (num_frames - 1) |
|
|
| if "beta" in motion_file_data: |
| _motion_aa.append(motion_file_data["pose_aa"].reshape(-1, self.num_joints * 3)) |
| _motion_bodies.append(curr_motion.gender_beta) |
| else: |
| _motion_aa.append(np.zeros((num_frames, self.num_joints * 3))) |
| _motion_bodies.append(torch.zeros(17)) |
|
|
| _motion_fps.append(motion_fps) |
| _motion_dt.append(curr_dt) |
| _motion_num_frames.append(num_frames) |
| motions.append(curr_motion) |
| _motion_lengths.append(curr_len) |
| if self.has_action: |
| _motion_actions.append(curr_motion.action) |
| if self.smpl_data is not None: |
| _motion_smpl_poses.append(curr_motion["smpl_pose"]) |
| if "smpl_joints" in curr_motion: |
| _motion_smpl_joints.append(curr_motion["smpl_joints"]) |
| if "smpl_transl" in curr_motion: |
| _motion_smpl_transl.append(curr_motion["smpl_transl"]) |
| if self.soma_data is not None: |
| if "soma_joints" in curr_motion: |
| _motion_soma_joints.append(curr_motion["soma_joints"]) |
| if "soma_root_quat" in curr_motion: |
| _motion_soma_root_quat.append(curr_motion["soma_root_quat"]) |
| if "soma_transl" in curr_motion: |
| _motion_soma_transl.append(curr_motion["soma_transl"]) |
| if self.object_data is not None: |
| if "object_root_pos" in curr_motion: |
| _motion_object_root_pos.append(curr_motion["object_root_pos"]) |
| if "object_root_quat" in curr_motion: |
| _motion_object_root_quat.append(curr_motion["object_root_quat"]) |
| if "object_contact_center_left" in curr_motion: |
| _motion_object_contact_center_left.append( |
| curr_motion["object_contact_center_left"] |
| ) |
| if "object_in_contact_left" in curr_motion: |
| _motion_object_in_contact_left.append(curr_motion["object_in_contact_left"]) |
| if "object_contact_center_right" in curr_motion: |
| _motion_object_contact_center_right.append( |
| curr_motion["object_contact_center_right"] |
| ) |
| if "object_in_contact_right" in curr_motion: |
| _motion_object_in_contact_right.append(curr_motion["object_in_contact_right"]) |
| if "hand_action_left" in motion_file_data: |
| raw_action = motion_file_data["hand_action_left"] |
| |
| src_len = len(raw_action) |
| if src_len != num_frames: |
| indices = np.round(np.linspace(0, src_len - 1, num_frames)).astype(int) |
| raw_action = raw_action[indices] |
| _motion_hand_action_left.append(raw_action) |
| if "hand_action_right" in motion_file_data: |
| raw_action = motion_file_data["hand_action_right"] |
| |
| src_len = len(raw_action) |
| if src_len != num_frames: |
| indices = np.round(np.linspace(0, src_len - 1, num_frames)).astype(int) |
| raw_action = raw_action[indices] |
| _motion_hand_action_right.append(raw_action) |
| del curr_motion |
|
|
| self._motion_lengths = torch.tensor( |
| _motion_lengths, device=self._device, dtype=torch.float32 |
| ) |
| self._motion_fps = torch.tensor(_motion_fps, device=self._device, dtype=torch.float32) |
| self._motion_bodies = torch.stack(_motion_bodies).to(self._device).type(torch.float32) |
| self._motion_aa = torch.tensor( |
| np.concatenate(_motion_aa), device=self._device, dtype=torch.float32 |
| ) |
|
|
| if self.smpl_data is not None: |
| self._motion_smpl_poses = torch.cat(_motion_smpl_poses, dim=0).float().to(self._device) |
| if len(_motion_smpl_joints) > 0: |
| self._motion_smpl_joints = ( |
| torch.cat(_motion_smpl_joints, dim=0).float().to(self._device) |
| ) |
| if len(_motion_smpl_transl) > 0: |
| self._motion_smpl_transl = ( |
| torch.cat(_motion_smpl_transl, dim=0).float().to(self._device) |
| ) |
| if self.soma_data is not None: |
| if len(_motion_soma_joints) > 0: |
| self._motion_soma_joints = ( |
| torch.cat(_motion_soma_joints, dim=0).float().to(self._device) |
| ) |
| if len(_motion_soma_root_quat) > 0: |
| self._motion_soma_root_quat = ( |
| torch.cat(_motion_soma_root_quat, dim=0).float().to(self._device) |
| ) |
| if len(_motion_soma_transl) > 0: |
| self._motion_soma_transl = ( |
| torch.cat(_motion_soma_transl, dim=0).float().to(self._device) |
| ) |
| if self.object_data is not None: |
| if len(_motion_object_root_pos) > 0: |
| self._motion_object_root_pos = ( |
| torch.cat(_motion_object_root_pos, dim=0).float().to(self._device) |
| ) |
| if len(_motion_object_root_quat) > 0: |
| self._motion_object_root_quat = ( |
| torch.cat(_motion_object_root_quat, dim=0).float().to(self._device) |
| ) |
| |
| if len(_motion_object_contact_center_left) > 0: |
| self._motion_object_contact_center_left = ( |
| torch.cat(_motion_object_contact_center_left, dim=0).float().to(self._device) |
| ) |
| if len(_motion_object_in_contact_left) > 0: |
| self._motion_object_in_contact_left = ( |
| torch.cat(_motion_object_in_contact_left, dim=0).float().to(self._device) |
| ) |
| if len(_motion_object_contact_center_right) > 0: |
| self._motion_object_contact_center_right = ( |
| torch.cat(_motion_object_contact_center_right, dim=0).float().to(self._device) |
| ) |
| if len(_motion_object_in_contact_right) > 0: |
| self._motion_object_in_contact_right = ( |
| torch.cat(_motion_object_in_contact_right, dim=0).float().to(self._device) |
| ) |
| if len(_motion_hand_action_left) > 0: |
| self._motion_hand_action_left = ( |
| torch.from_numpy(np.concatenate(_motion_hand_action_left, axis=0)) |
| .float() |
| .to(self._device) |
| ) |
| if len(_motion_hand_action_right) > 0: |
| self._motion_hand_action_right = ( |
| torch.from_numpy(np.concatenate(_motion_hand_action_right, axis=0)) |
| .float() |
| .to(self._device) |
| ) |
| self._motion_dt = torch.tensor(_motion_dt, device=self._device, dtype=torch.float32) |
|
|
| |
| if self.object_data is not None and hasattr(self, "_motion_object_root_pos"): |
| self._compute_object_velocities(_motion_num_frames, _motion_dt) |
| self._motion_num_frames = torch.tensor(_motion_num_frames, device=self._device) |
|
|
| if self.has_action: |
| self._motion_actions = torch.cat(_motion_actions, dim=0).float().to(self._device) |
| self._num_motions = len(motions) |
|
|
| self.body_pos_w = ( |
| torch.cat([m.global_translation for m in motions], dim=0).float().to(self._device) |
| ) |
| self.body_quat_w = ( |
| torch.cat([m.global_rotation for m in motions], dim=0).float().to(self._device) |
| ) |
| self.body_pos_b = ( |
| torch.cat([m.local_rotation for m in motions], dim=0).float().to(self._device) |
| ) |
| self.root_linv_vel_w = ( |
| torch.cat([m.global_root_velocity for m in motions], dim=0).float().to(self._device) |
| ) |
| self.root_ang_vel_w = ( |
| torch.cat([m.global_root_angular_velocity for m in motions], dim=0) |
| .float() |
| .to(self._device) |
| ) |
| self.body_ang_vel_w = ( |
| torch.cat([m.global_angular_velocity for m in motions], dim=0).float().to(self._device) |
| ) |
| self.body_lin_vel_w = ( |
| torch.cat([m.global_velocity for m in motions], dim=0).float().to(self._device) |
| ) |
| self.dof_vel = torch.cat([m.dof_vels for m in motions], dim=0).float().to(self._device) |
| self.feet_l = torch.cat([m.feet_l for m in motions], dim=0).float().to(self._device) |
| self.feet_r = torch.cat([m.feet_r for m in motions], dim=0).float().to(self._device) |
|
|
| |
| |
| |
| |
| |
| |
| |
|
|
| if "dof_pos" in motions[0].__dict__: |
| self.dof_pos = torch.cat([m.dof_pos for m in motions], dim=0).float().to(self._device) |
|
|
| |
| if "hand_dof_pos" in motions[0].__dict__: |
| self.hand_dof_pos = ( |
| torch.cat([m.hand_dof_pos for m in motions], dim=0).float().to(self._device) |
| ) |
| else: |
| self.hand_dof_pos = None |
|
|
| lengths = self._motion_num_frames |
| lengths_shifted = lengths.roll(1) |
| lengths_shifted[0] = 0 |
| self.length_starts = lengths_shifted.cumsum(0) |
|
|
| |
| if self.m_cfg.get("zero_root_xy", False): |
| print( |
| f"[zero_root_xy] Zeroing initial root XY for {len(motions)} motions" |
| ) |
| for i in range(len(motions)): |
| start = self.length_starts[i] |
| end = start + self._motion_num_frames[i] |
| init_xy = self.body_pos_w[start, 0, :2].clone() |
| print( |
| f" Motion {i}: init_xy=[{init_xy[0]:.3f}, {init_xy[1]:.3f}], frames={self._motion_num_frames[i]}" |
| ) |
| self.body_pos_w[start:end, :, :2] -= init_xy |
| if ( |
| hasattr(self, "_motion_object_root_pos") |
| and self._motion_object_root_pos is not None |
| ): |
| self._motion_object_root_pos[start:end, :, :2] -= init_xy |
|
|
| self.motion_ids = torch.arange(len(motions), dtype=torch.long, device=self._device) |
|
|
| motion_has_smpl = [ |
| self.curr_motion_keys[i] in self.smpl_data_keys for i in range(len(motions)) |
| ] |
| self.motion_has_smpl = torch.tensor(motion_has_smpl, dtype=torch.bool, device=self._device) |
|
|
| motion_has_soma = [ |
| self.curr_motion_keys[i] in self.soma_data_keys for i in range(len(motions)) |
| ] |
| self.motion_has_soma = torch.tensor(motion_has_soma, dtype=torch.bool, device=self._device) |
|
|
| motion_has_object = [ |
| self.curr_motion_keys[i] in self.object_data_keys for i in range(len(motions)) |
| ] |
| self.motion_has_object = torch.tensor( |
| motion_has_object, dtype=torch.bool, device=self._device |
| ) |
|
|
| motion = motions[0] |
| self.num_bodies = self.num_joints |
|
|
| num_motions = self.num_motions() |
| total_len = self.get_total_length() |
|
|
| if self.use_adaptive_sampling: |
| self.update_adaptive_sampling_motion_frames() |
|
|
| logger.info( |
| f"Loaded {num_motions:d} motions with a total length of {total_len:.3f}s and {self.body_pos_w.shape[0]} frames." |
| ) |
|
|
| del ( |
| motions, |
| _motion_lengths, |
| _motion_fps, |
| _motion_dt, |
| _motion_num_frames, |
| _motion_bodies, |
| _motion_aa, |
| _motion_actions, |
| _motion_smpl_poses, |
| _motion_smpl_joints, |
| _motion_smpl_transl, |
| _motion_object_root_pos, |
| _motion_object_root_quat, |
| ) |
|
|
| if "mujoco_to_isaaclab_body" in self.m_cfg.keys(): |
| self.dof_pos = self.dof_pos[:, self.m_cfg.mujoco_to_isaaclab_dof] |
| self.dof_vel = self.dof_vel[:, self.m_cfg.mujoco_to_isaaclab_dof] |
|
|
| |
| self.num_bodies_full = len(self.m_cfg.mujoco_to_isaaclab_body) |
| self.body_pos_w_full = self.body_pos_w[:, self.m_cfg.mujoco_to_isaaclab_body] |
| self.body_quat_w_full = rotations.xyzw_to_wxyz( |
| self.body_quat_w[:, self.m_cfg.mujoco_to_isaaclab_body] |
| ) |
| self.body_lin_vel_w_full = self.body_lin_vel_w[:, self.m_cfg.mujoco_to_isaaclab_body] |
| self.body_ang_vel_w_full = self.body_ang_vel_w[:, self.m_cfg.mujoco_to_isaaclab_body] |
|
|
| |
| self.body_pos_w = self.body_pos_w_full[:, self.body_indexes] |
| self.body_quat_w = self.body_quat_w_full[:, self.body_indexes] |
| self.body_lin_vel_w = self.body_lin_vel_w_full[:, self.body_indexes] |
| self.body_ang_vel_w = self.body_ang_vel_w_full[:, self.body_indexes] |
| assert ( |
| self.m_cfg.get("anchor_body_idx_full", 0) == 0 and self.body_indexes[0] == 0 |
| ), "The anchor body has to be 0; otherwise will cause issues in the sliced body_indexes data's anchor." |
| else: |
| |
| self.body_pos_w_full = self.body_pos_w |
| self.body_quat_w_full = self.body_quat_w |
| self.body_lin_vel_w_full = self.body_lin_vel_w |
| self.body_ang_vel_w_full = self.body_ang_vel_w |
| self.num_bodies_full = self.body_pos_w.shape[2] |
|
|
| |
| gc.collect() |
| torch.cuda.empty_cache() |
|
|
| def foot_detect(self, positions, vel_thres, height_thresh): |
| fid_l = self.m_cfg.get("left_foot_body_idx", [6]) |
| fid_r = self.m_cfg.get("right_foot_body_idx", [12]) |
| |
| velfactor = torch.tensor( |
| [vel_thres] * len(fid_l), device=positions.device, dtype=positions.dtype |
| ) |
| heightfactor = torch.tensor( |
| [height_thresh] * len(fid_l), device=positions.device, dtype=positions.dtype |
| ) |
|
|
| feet_l_xyz = (positions[1:, fid_l] - positions[:-1, fid_l]) ** 2 |
| feet_l_xyz = torch.cat([feet_l_xyz, feet_l_xyz[[-1]]], dim=0) |
| feet_l_h = positions[:, fid_l, 2] |
| feet_l = torch.logical_and( |
| (feet_l_xyz.sum(dim=-1)) < velfactor, feet_l_h < heightfactor |
| ).float() |
| |
|
|
| feet_r_xyz = (positions[1:, fid_r] - positions[:-1, fid_r]) ** 2 |
| feet_r_xyz = torch.cat([feet_r_xyz, feet_r_xyz[[-1]]], dim=0) |
| feet_r_h = positions[:, fid_r, 2] |
| feet_r = torch.logical_and( |
| (feet_r_xyz.sum(dim=-1)) < velfactor, feet_r_h < heightfactor |
| ).float() |
| |
| return feet_l, feet_r |
|
|
| def _compute_object_velocities(self, motion_num_frames, motion_dt): |
| """Compute object linear and angular velocities from position and quaternion data. |
| Uses finite differences: v = (p_{t+1} - p_t) / dt |
| Handles motion boundaries properly (first frame uses forward difference). |
| """ |
| total_frames = self._motion_object_root_pos.shape[0] |
| num_objects = self._motion_object_root_pos.shape[1] |
|
|
| |
| self._motion_object_lin_vel = torch.zeros_like(self._motion_object_root_pos) |
| self._motion_object_ang_vel = torch.zeros( |
| total_frames, num_objects, 3, device=self._device, dtype=torch.float32 |
| ) |
|
|
| |
| num_frames_tensor = torch.tensor(motion_num_frames, device=self._device) |
| lengths_shifted = num_frames_tensor.roll(1) |
| lengths_shifted[0] = 0 |
| length_starts = lengths_shifted.cumsum(0) |
|
|
| |
| for i, (start, num_frames, dt) in enumerate( |
| zip(length_starts, motion_num_frames, motion_dt) |
| ): |
| start = start.item() |
| end = start + num_frames |
|
|
| if num_frames < 2: |
| continue |
|
|
| |
| pos = self._motion_object_root_pos[start:end] |
| quat = self._motion_object_root_quat[start:end] |
|
|
| |
| lin_vel = (pos[1:] - pos[:-1]) / dt |
| |
| lin_vel = torch.cat([lin_vel[:1], lin_vel], dim=0) |
| self._motion_object_lin_vel[start:end] = lin_vel |
|
|
| |
| |
| q_curr = quat[:-1] |
| q_next = quat[1:] |
|
|
| |
| |
| diff_quat = rotations.quat_mul_norm( |
| q_next, rotations.quat_inverse(q_curr, w_last=False), w_last=False |
| ) |
|
|
| |
| diff_angle, diff_axis = rotations.quat_angle_axis(diff_quat, w_last=False) |
|
|
| |
| ang_vel = diff_axis * diff_angle.unsqueeze(-1) / dt |
| |
| ang_vel = torch.cat([ang_vel[:1], ang_vel], dim=0) |
| self._motion_object_ang_vel[start:end] = ang_vel |
|
|
| logger.info(f"Computed object velocities for {len(motion_num_frames)} motions") |
|
|
| def fix_trans_height(self, pose_aa, trans, fix_height_mode): |
| if fix_height_mode == FixHeightMode.no_fix: |
| return trans, 0 |
| with torch.no_grad(): |
|
|
| mesh_obj = self.mesh_parsers.mesh_fk(pose_aa[None, :1], trans[None, :1]) |
| height_diff = np.asarray(mesh_obj.vertices)[..., 2].min() |
| trans[..., 2] -= height_diff |
|
|
| return trans, height_diff |
|
|
| def load_motion_with_skeleton( |
| self, |
| ids, |
| motion_data_list, |
| smpl_data_list, |
| object_data_list, |
| soma_data_list, |
| fix_height, |
| target_heading, |
| max_len, |
| is_evaluation, |
| queue, |
| pid, |
| ): |
| |
| res = {} |
|
|
| if pid == 0: |
| pbar = progress.track(range(len(ids)), description="Loading motions...") |
| else: |
| pbar = range(len(ids)) |
|
|
| for f in pbar: |
|
|
| curr_id = ids[f] |
|
|
| curr_file = motion_data_list[f] |
| if "path" in curr_file: |
| curr_file, *_ = joblib.load( |
| curr_file["path"] |
| ).values() |
|
|
| seq_len = curr_file["root_trans_offset"].shape[0] |
| if max_len == -1 or seq_len < max_len: |
| start, end = 0, seq_len |
| else: |
| start = random.randint(0, seq_len - max_len) |
| end = start + max_len |
|
|
| trans = to_torch(curr_file["root_trans_offset"]).clone()[start:end] |
| pose_aa = to_torch(curr_file["pose_aa"][start:end]).clone() |
|
|
| |
| if "action" in curr_file.keys(): |
| self.has_action = True |
|
|
| if "fps" not in curr_file.keys(): |
| curr_file["fps"] = 30.0 |
| dt = 1 / curr_file["fps"] |
|
|
| B, J, N = pose_aa.shape |
| freeze_frame_aug, freeze_idx = False, 0 |
|
|
| |
| |
|
|
| if not is_evaluation and self.m_cfg.get("freeze_frame_aug", False): |
| freeze_prob = self.m_cfg.get("freeze_frame_prob", 0.1) |
| if np.random.random() < freeze_prob: |
| |
| freeze_frame_aug = True |
| freeze_idx = np.random.randint(0, B) |
| |
| pose_aa[freeze_idx:] = pose_aa[freeze_idx : freeze_idx + 1].clone() |
| trans[freeze_idx:] = trans[freeze_idx : freeze_idx + 1].clone() |
|
|
| if not is_evaluation and self.m_cfg.get("randomize_heading", False): |
| |
| random_rot = np.zeros(3) |
| random_rot[2] = np.pi * (2 * np.random.random() - 1.0) |
| random_heading_rot = transform.Rotation.from_euler("xyz", random_rot) |
| pose_aa = pose_aa.reshape(B, -1) |
| pose_aa[:, :3] = torch.tensor( |
| ( |
| random_heading_rot * transform.Rotation.from_rotvec(pose_aa[:, :3]) |
| ).as_rotvec() |
| ) |
| trans = torch.matmul( |
| trans, torch.from_numpy(random_heading_rot.as_matrix().T).float() |
| ) |
| pose_aa = pose_aa.reshape(B, J, N) |
|
|
| |
| randomize_upper_body_poses = ( |
| self.randomize_upper_body_poses |
| and random.random() < self.cat_upper_body_poses_prob |
| and (self._should_augment_upper_body(self.curr_motion_keys[curr_id])) |
| ) |
|
|
| |
| if ( |
| randomize_upper_body_poses |
| and self.res_non_nav_dataset is not None |
| and len(self.res_non_nav_dataset) > 0 |
| ): |
| |
| |
| upper_body_indices = [ |
| i for i in range(1, J) if i - 1 not in self.m_cfg.lower_joint_indices_mujoco |
| ] |
| |
| selected_file, selected_motion = random.choice( |
| list(self.res_non_nav_dataset.values()) |
| ) |
| selected_pose_aa = to_torch(selected_file["pose_aa"]) |
|
|
| |
| |
| selected_seq_len = selected_file["root_trans_offset"].shape[0] |
| current_seq_len = pose_aa.shape[0] |
| if selected_seq_len >= current_seq_len: |
| selected_start = random.randint(0, selected_seq_len - current_seq_len) |
| selected_end = selected_start + current_seq_len |
| selected_slice = selected_pose_aa[selected_start:selected_end] |
| else: |
| |
| forward = selected_pose_aa |
| backward = selected_pose_aa.flip(dims=[0]) |
| |
| extended = torch.cat([forward, backward[1:]], dim=0) |
|
|
| |
| if extended.shape[0] < current_seq_len: |
| repeats = (current_seq_len + extended.shape[0] - 1) // extended.shape[ |
| 0 |
| ] |
| extended = extended.repeat(repeats, 1, 1) |
|
|
| selected_slice = extended[:current_seq_len] |
|
|
| pose_aa[:, upper_body_indices] = selected_slice[:, upper_body_indices] |
|
|
| |
| if ( |
| not is_evaluation |
| and self.randomize_wrist_poses |
| and random.random() < self.randomize_wrist_prob |
| ): |
| wrist_pose_aa_indices = [d + 1 for d in self.wrist_mujoco_dof_indices] |
| noise = torch.randn(B, len(wrist_pose_aa_indices), N) * self.randomize_wrist_std |
| pose_aa[:, wrist_pose_aa_indices] = pose_aa[:, wrist_pose_aa_indices] + noise |
|
|
| if self.mesh_parsers is not None: |
| trans, trans_fix = self.fix_trans_height(pose_aa, trans, fix_height_mode=fix_height) |
| curr_motion = self.mesh_parsers.fk_batch( |
| pose_aa[None,], |
| trans[None,], |
| return_full=True, |
| fps=curr_file["fps"], |
| target_fps=self.target_fps, |
| interpolate_data=True, |
| use_parallel_fk=self.use_parallel_fk, |
| ) |
| if self.smpl_data is not None: |
| curr_smpl_data = smpl_data_list[f] |
| if curr_smpl_data is not None: |
| if "path" in curr_smpl_data: |
| curr_smpl_data = joblib.load(curr_smpl_data["path"]) |
|
|
| if curr_smpl_data["fps"] != self.target_fps: |
| smpl_pose = torch.tensor(curr_smpl_data["pose_aa"][start:end]).float() |
| smpl_pose[:, -6:] = 0.0 |
| curr_motion["smpl_pose"] = self.mesh_parsers.interploate_pose( |
| None, smpl_pose[None,], curr_smpl_data["fps"], self.target_fps |
| )[1][0] |
| else: |
| smpl_pose = torch.tensor(curr_smpl_data["pose_aa"]).float() |
| smpl_pose[:, -6:] = 0.0 |
| |
| curr_motion["smpl_pose"] = smpl_pose |
| if "smpl_joints" in curr_smpl_data: |
| smpl_joints = torch.tensor(curr_smpl_data["smpl_joints"]).float() |
| curr_motion["smpl_joints"] = smpl_joints |
| if ( |
| curr_motion["smpl_joints"].shape[0] |
| != curr_motion["global_translation"].shape[1] |
| ): |
| print( |
| f"Length mismatch: smpl_joints={curr_motion['smpl_joints'].shape[0]}, " |
| f"global_translation={curr_motion['global_translation'].shape[1]}" |
| ) |
| print(smpl_data_list[f], motion_data_list[f]) |
|
|
| assert ( |
| curr_motion["smpl_joints"].shape[0] |
| == curr_motion["global_translation"].shape[1] |
| ) |
| else: |
| num_frames = curr_motion["global_translation"].shape[1] |
| curr_motion["smpl_joints"] = torch.zeros(num_frames, 24, 3).to( |
| curr_motion["global_translation"] |
| ) |
| if "transl" in curr_smpl_data: |
| transl = torch.tensor(curr_smpl_data["transl"]).float() |
| curr_motion["smpl_transl"] = transl |
| assert ( |
| curr_motion["smpl_transl"].shape[0] |
| == curr_motion["global_translation"].shape[1] |
| ) |
| else: |
| num_frames = curr_motion["global_translation"].shape[1] |
| curr_motion["smpl_transl"] = torch.zeros(num_frames, 3).to( |
| curr_motion["global_translation"] |
| ) |
| assert ( |
| curr_motion["smpl_pose"].shape[0] |
| == curr_motion["global_translation"].shape[1] |
| ) |
|
|
| if freeze_frame_aug: |
| freeze_idx_new_fps = int( |
| freeze_idx * self.target_fps / curr_file["fps"] |
| ) |
| curr_motion["smpl_pose"][freeze_idx_new_fps:] = curr_motion[ |
| "smpl_pose" |
| ][freeze_idx_new_fps : freeze_idx_new_fps + 1].clone() |
| curr_motion["smpl_joints"][freeze_idx_new_fps:] = curr_motion[ |
| "smpl_joints" |
| ][freeze_idx_new_fps : freeze_idx_new_fps + 1].clone() |
| curr_motion["smpl_transl"][freeze_idx_new_fps:] = curr_motion[ |
| "smpl_transl" |
| ][freeze_idx_new_fps : freeze_idx_new_fps + 1].clone() |
| else: |
| curr_motion["smpl_pose"] = torch.zeros( |
| curr_motion["global_translation"].shape[1], 72 |
| ).to(curr_motion["global_translation"]) |
| curr_motion["smpl_joints"] = torch.zeros( |
| curr_motion["global_translation"].shape[1], 24, 3 |
| ).to(curr_motion["global_translation"]) |
| curr_motion["smpl_transl"] = torch.zeros( |
| curr_motion["global_translation"].shape[1], 3 |
| ).to(curr_motion["global_translation"]) |
| |
|
|
| |
| if soma_data_list is not None: |
| curr_soma_data = soma_data_list[f] |
| if curr_soma_data is not None: |
| if "path" in curr_soma_data: |
| loaded = joblib.load(curr_soma_data["path"]) |
| curr_soma_data, *_ = loaded.values() |
|
|
| num_frames = curr_motion["global_translation"].shape[1] |
| n_soma = self.num_soma_joints |
|
|
| |
| |
| soma_fps = curr_soma_data.get("fps", self.target_fps) |
|
|
| if "soma_joints" in curr_soma_data: |
| soma_joints = torch.tensor(curr_soma_data["soma_joints"]).float() |
| soma_joints_orig_len = soma_joints.shape[0] |
| if soma_fps != self.target_fps: |
| soma_joints = self._resample_soma_tensor( |
| soma_joints, soma_fps, self.target_fps |
| ) |
| curr_motion["soma_joints"] = soma_joints |
| assert soma_joints.shape[0] == num_frames, ( |
| f"SOMA soma_joints length {soma_joints.shape[0]} != " |
| f"robot frames {num_frames} " |
| f"(soma_orig={soma_joints_orig_len} @ {soma_fps}fps, " |
| f"robot_orig={seq_len} @ {curr_file['fps']}fps, " |
| f"target_fps={self.target_fps})" |
| ) |
| else: |
| curr_motion["soma_joints"] = torch.zeros(num_frames, n_soma, 3).to( |
| curr_motion["global_translation"] |
| ) |
|
|
| if "soma_root_quat" in curr_soma_data: |
| soma_root_quat = torch.tensor(curr_soma_data["soma_root_quat"]).float() |
| if soma_fps != self.target_fps: |
| soma_root_quat = self._resample_soma_tensor( |
| soma_root_quat, soma_fps, self.target_fps |
| ) |
| |
| soma_root_quat = soma_root_quat / ( |
| soma_root_quat.norm(dim=-1, keepdim=True) + 1e-8 |
| ) |
| curr_motion["soma_root_quat"] = soma_root_quat |
| assert soma_root_quat.shape[0] == num_frames, ( |
| f"SOMA soma_root_quat length {soma_root_quat.shape[0]} != " |
| f"robot frames {num_frames}" |
| ) |
| else: |
| curr_motion["soma_root_quat"] = torch.zeros(num_frames, 4).to( |
| curr_motion["global_translation"] |
| ) |
| curr_motion["soma_root_quat"][:, 0] = 1.0 |
|
|
| if "soma_transl" in curr_soma_data: |
| soma_transl = torch.tensor(curr_soma_data["soma_transl"]).float() |
| if soma_fps != self.target_fps: |
| soma_transl = self._resample_soma_tensor( |
| soma_transl, soma_fps, self.target_fps |
| ) |
| curr_motion["soma_transl"] = soma_transl |
| assert soma_transl.shape[0] == num_frames, ( |
| f"SOMA soma_transl length {soma_transl.shape[0]} != " |
| f"robot frames {num_frames}" |
| ) |
| else: |
| curr_motion["soma_transl"] = torch.zeros(num_frames, 3).to( |
| curr_motion["global_translation"] |
| ) |
|
|
| if freeze_frame_aug: |
| freeze_idx_new_fps = int( |
| freeze_idx * self.target_fps / curr_file["fps"] |
| ) |
| for key in ("soma_joints", "soma_root_quat", "soma_transl"): |
| curr_motion[key][freeze_idx_new_fps:] = curr_motion[key][ |
| freeze_idx_new_fps : freeze_idx_new_fps + 1 |
| ].clone() |
| else: |
| num_frames = curr_motion["global_translation"].shape[1] |
| n_soma = self.num_soma_joints |
| curr_motion["soma_joints"] = torch.zeros(num_frames, n_soma, 3).to( |
| curr_motion["global_translation"] |
| ) |
| curr_motion["soma_root_quat"] = torch.zeros(num_frames, 4).to( |
| curr_motion["global_translation"] |
| ) |
| curr_motion["soma_root_quat"][:, 0] = 1.0 |
| curr_motion["soma_transl"] = torch.zeros(num_frames, 3).to( |
| curr_motion["global_translation"] |
| ) |
|
|
| |
| if self.object_data is not None: |
| curr_object_data = object_data_list[f] if object_data_list is not None else None |
| if curr_object_data is not None: |
| if "path" in curr_object_data: |
| loaded = joblib.load(curr_object_data["path"]) |
| curr_object_data, *_ = loaded.values() |
|
|
| num_frames = curr_motion["global_translation"].shape[1] |
| original_fps = curr_object_data.get("fps", curr_file["fps"]) |
|
|
| if "root_pos" in curr_object_data: |
| curr_motion["object_root_pos"] = interpolate_translation_data( |
| curr_object_data["root_pos"][start:end], |
| source_fps=original_fps, |
| target_fps=self.target_fps, |
| num_frames=num_frames, |
| max_num_objects=self.max_num_objects, |
| pad_value=0.0, |
| ) |
|
|
| if "root_quat" in curr_object_data: |
| curr_motion["object_root_quat"] = interpolate_quaternion_data( |
| curr_object_data["root_quat"][start:end], |
| source_fps=original_fps, |
| target_fps=self.target_fps, |
| num_frames=num_frames, |
| max_num_objects=self.max_num_objects, |
| ) |
|
|
| |
| for hand in ("left_hand", "right_hand"): |
| key = f"contact_points_{hand}" |
| side = hand.split("_")[0] |
| if key in curr_object_data: |
| |
| |
| raw_dict = curr_object_data[key] |
| sliced_dict = { |
| k - start: v for k, v in raw_dict.items() if start <= k < end |
| } |
| center, label = interpolate_contact_center( |
| sliced_dict, |
| source_fps=original_fps, |
| target_fps=self.target_fps, |
| num_frames=num_frames, |
| ) |
| curr_motion[f"object_contact_center_{side}"] = center |
| curr_motion[f"object_in_contact_{side}"] = label |
| else: |
| |
| num_frames = curr_motion["global_translation"].shape[1] |
| curr_motion["object_root_pos"] = torch.zeros( |
| num_frames, self.max_num_objects, 3 |
| ).to(curr_motion["global_translation"]) |
| curr_motion["object_root_quat"] = torch.zeros( |
| num_frames, self.max_num_objects, 4 |
| ).to(curr_motion["global_translation"]) |
| curr_motion["object_root_quat"][ |
| :, :, 0 |
| ] = 1.0 |
|
|
| curr_motion = easydict.EasyDict( |
| { |
| k: v.squeeze(dim=-1).squeeze(dim=0) if torch.is_tensor(v) else v |
| for k, v in curr_motion.items() |
| } |
| ) |
| |
| if self.has_action: |
| curr_motion.action = to_torch(curr_file["action"]).clone()[start:end] |
|
|
| |
| hand_dof_count = self.m_cfg.get("hand_dof_count", 0) |
| if hand_dof_count > 0 and "dof" in curr_file: |
| raw_dof = to_torch(curr_file["dof"]).clone()[start:end] |
| if raw_dof.shape[-1] > 29: |
| |
| hand_dof = raw_dof[:, 29 : 29 + hand_dof_count] |
| if curr_file["fps"] != self.target_fps: |
| |
| num_target_frames = curr_motion["dof_pos"].shape[0] |
| hand_dof_interp = ( |
| torch.nn.functional.interpolate( |
| hand_dof.T.unsqueeze(0), |
| size=num_target_frames, |
| mode="linear", |
| align_corners=True, |
| ) |
| .squeeze(0) |
| .T |
| ) |
| curr_motion.hand_dof_pos = hand_dof_interp |
| else: |
| curr_motion.hand_dof_pos = hand_dof |
|
|
| if self.vid_smpl_pose is not None: |
| vid_smpl_pose = self.vid_smpl_pose[f] |
| vid_smpl_pose = self.mesh_parsers.interploate_pose( |
| None, vid_smpl_pose[None,], 30.0, self.target_fps |
| )[1][0] |
| if curr_motion["smpl_pose"].shape[0] < vid_smpl_pose.shape[0]: |
| for key in curr_motion.keys(): |
| if isinstance(curr_motion[key], torch.Tensor): |
| curr_motion[key] = torch.cat( |
| [ |
| curr_motion[key], |
| torch.zeros( |
| vid_smpl_pose.shape[0] - curr_motion[key].shape[0], |
| *curr_motion[key].shape[1:], |
| ).to(curr_motion[key]), |
| ], |
| dim=0, |
| ) |
| curr_motion["smpl_pose"][:, :] = vid_smpl_pose[ |
| : curr_motion["smpl_pose"].shape[0], : |
| ] |
| feet_l, feet_r = self.foot_detect(curr_motion["global_translation"], 0.0005, 0.05) |
| curr_motion["feet_l"] = feet_l |
| curr_motion["feet_r"] = feet_r |
| res[curr_id] = (curr_file, curr_motion) |
| else: |
| logger.error("No mesh parser found") |
|
|
| if queue is not None: |
| queue.put(res) |
| else: |
| return res |
|
|
| def num_motions(self): |
| return self._num_motions |
|
|
| def get_total_length(self): |
| return sum(self._motion_lengths) |
|
|
| def get_motion_num_steps(self, motion_ids=None): |
| if motion_ids is None: |
| return ( |
| (self._motion_num_frames * self._sim_fps / self._motion_fps).floor().int() |
| ) |
| else: |
| return ( |
| (self._motion_num_frames[motion_ids] * self._sim_fps / self._motion_fps[motion_ids]) |
| .floor() |
| .int() |
| ) |
|
|
| def sample_time(self, motion_ids, truncate_time=None): |
| n = len(motion_ids) |
| phase = torch.rand(motion_ids.shape, device=self._device) |
| motion_len = self._motion_lengths[motion_ids] |
| if truncate_time is not None: |
| assert truncate_time >= 0.0 |
| motion_len -= truncate_time |
|
|
| motion_time = phase * motion_len |
| return motion_time.to(self._device) |
|
|
| def sample_time_steps(self, motion_ids, truncate_time=None): |
| motion_time = self.sample_time(motion_ids, truncate_time) |
| motion_time_steps = (motion_time * self._sim_fps).floor().int() |
| return motion_time_steps |
|
|
| def sample_motions(self, n): |
| motion_ids = torch.multinomial( |
| self._sampling_batch_prob, num_samples=n, replacement=True |
| ).to(self._device) |
|
|
| return motion_ids |
|
|
| def get_motion_ids_in_dataset(self, motion_ids): |
| return self._curr_motion_ids[motion_ids] |
|
|
| def get_motion_length(self, motion_ids=None): |
| if motion_ids is None: |
| return self._motion_lengths |
| else: |
| return self._motion_lengths[motion_ids] |
|
|
| def _calc_frame_blend(self, time, len, num_frames, dt): |
| time = time.clone() |
| phase = time / len |
| phase = torch.clip(phase, 0.0, 1.0) |
| time[time < 0] = 0 |
|
|
| frame_idx0 = (phase * (num_frames - 1)).long() |
| frame_idx1 = torch.min(frame_idx0 + 1, num_frames - 1) |
| blend = torch.clip( |
| (time - frame_idx0 * dt) / dt, 0.0, 1.0 |
| ) |
|
|
| return frame_idx0, frame_idx1, blend |
|
|
| def _get_num_bodies(self): |
| return self.num_bodies |
|
|
| def _local_rotation_to_dof_smpl(self, local_rot): |
| B, J, _ = local_rot.shape |
| dof_pos = rotations.quat_to_exp_map(local_rot[:, 1:]) |
| return dof_pos.reshape(B, -1) |
|
|
| def init_adaptive_sampling(self): |
| """Initialize adaptive sampling data structures over all unique motions. |
| |
| Divides every motion clip into fixed-size bins (``bin_size`` frames each) and |
| creates per-bin tracking tensors for failure rates and sampling probabilities. |
| This enables fine-grained, time-segment-level curriculum learning: bins with |
| higher failure rates are sampled more frequently during training. |
| |
| NOTE: This operates over ALL unique motions in the dataset (not just the |
| currently loaded batch), so bin indices are stable across reloads. |
| """ |
| self.adp_samp_num_frames = torch.zeros( |
| self._num_unique_motions, device=self._device, dtype=torch.long |
| ) |
| |
| for i, motion_key in enumerate(self._motion_data_keys): |
| motion_data = self._motion_data_load[motion_key] |
|
|
| |
| |
| if "fps" not in motion_data.keys(): |
| motion_data["fps"] = 30.0 |
|
|
| original_fps = motion_data["fps"] |
| |
| |
| if "length" in motion_data: |
| original_num_frames = motion_data["length"] |
| elif "root_trans_offset" in motion_data: |
| original_num_frames = motion_data["root_trans_offset"].shape[0] |
| elif "path" in motion_data: |
| |
| loaded_data, *_ = joblib.load(motion_data["path"]).values() |
| original_num_frames = loaded_data["root_trans_offset"].shape[0] |
| if "fps" in loaded_data: |
| original_fps = loaded_data["fps"] |
| else: |
| raise KeyError( |
| f"Cannot determine frame count for motion '{motion_key}': no 'length', 'root_trans_offset', or 'path' key" |
| ) |
|
|
| original_duration = (original_num_frames - 1) / original_fps |
|
|
| |
| |
| |
| if original_fps == self.target_fps: |
| num_frames = original_num_frames |
| else: |
| num_frames = len(torch.arange(0, original_duration, 1 / self.target_fps)) |
| self.adp_samp_num_frames[i] = num_frames |
|
|
| |
| lengths = self.adp_samp_num_frames |
| lengths_shifted = lengths.roll(1) |
| lengths_shifted[0] = 0 |
| self.adp_samp_length_starts = lengths_shifted.cumsum(0) |
| self.adp_samp_total_frames = self.adp_samp_num_frames.sum() |
| self.adp_samp_length_starts_mask = torch.zeros( |
| self.adp_samp_total_frames, device=self._device, dtype=torch.bool |
| ) |
| self.adp_samp_length_starts_mask[self.adp_samp_length_starts] = True |
|
|
| |
| self.adp_samp_bin_size = self.adaptive_sampling_cfg.get("bin_size", 50) |
| self.adp_samp_frame_to_bin = torch.zeros( |
| self.adp_samp_total_frames, device=self._device, dtype=torch.long |
| ) |
|
|
| |
| all_bins = [] |
| all_bin_motion_lengths = [] |
| all_bin_new_motion_masks = [] |
| all_num_peer_bins = [] |
| all_motion_to_bins = [] |
|
|
| cur_bin_idx = 0 |
| for orig_motion_id in range(self._num_unique_motions): |
| num_frames = self.adp_samp_num_frames[orig_motion_id] |
| frame_start = self.adp_samp_length_starts[orig_motion_id] |
| frame_end = ( |
| self.adp_samp_length_starts[orig_motion_id + 1] |
| if orig_motion_id < self._num_unique_motions - 1 |
| else self.adp_samp_total_frames |
| ) |
|
|
| |
| bin_starts = torch.arange( |
| 0, num_frames, self.adp_samp_bin_size, device=self._device, dtype=torch.long |
| ) |
| bin_ends = torch.minimum(bin_starts + self.adp_samp_bin_size, num_frames) |
| num_bins = len(bin_starts) |
| motion_ids = torch.full( |
| (num_bins,), orig_motion_id, device=self._device, dtype=torch.long |
| ) |
| motion_bins = torch.stack([motion_ids, bin_starts, bin_ends], dim=1) |
| all_bins.append(motion_bins) |
|
|
| |
| bin_lengths = bin_ends - bin_starts |
| all_bin_motion_lengths.append(bin_lengths) |
|
|
| |
| new_motion_mask = torch.zeros(num_bins, device=self._device, dtype=torch.bool) |
| new_motion_mask[0] = True |
| all_bin_new_motion_masks.append(new_motion_mask) |
|
|
| |
| peer_bins = torch.full((num_bins,), num_bins, device=self._device, dtype=torch.long) |
| all_num_peer_bins.append(peer_bins) |
|
|
| bin_ids = torch.zeros(num_frames, device=self._device, dtype=torch.long) |
| bin_ids[bin_starts[1:]] = 1 |
| bin_ids = bin_ids.cumsum(0) + cur_bin_idx |
| self.adp_samp_frame_to_bin[frame_start:frame_end] = bin_ids |
|
|
| |
| motion_bin_indices = torch.arange( |
| cur_bin_idx, cur_bin_idx + num_bins, device=self._device, dtype=torch.long |
| ) |
| all_motion_to_bins.append(motion_bin_indices) |
|
|
| cur_bin_idx += num_bins |
|
|
| |
| self.adp_samp_bins = torch.cat(all_bins, dim=0) |
| self.adp_samp_bin_motion_length = torch.cat(all_bin_motion_lengths, dim=0) |
| self.adp_samp_bin_new_motion_mask = torch.cat(all_bin_new_motion_masks, dim=0) |
| self.adp_samp_num_peer_bins = torch.cat(all_num_peer_bins, dim=0) |
| self.orig_motion_id_to_bins = all_motion_to_bins |
| self.adp_samp_num_bins = len(self.adp_samp_bins) |
|
|
| self.adp_samp_bin_weights = ( |
| self.adp_samp_bin_motion_length / self.adp_samp_bin_motion_length.float().mean() |
| ) |
| |
| if self.adaptive_sampling_cfg.get("sequence_length_agnostic", True): |
| self.adp_samp_bin_weights = self.adp_samp_bin_weights / self.adp_samp_num_peer_bins |
|
|
| init_num_failures = self.adaptive_sampling_cfg.get("init_num_failures", 1) |
| self.adp_samp_failure_rate_max_over_mean = self.adaptive_sampling_cfg.get( |
| "adp_samp_failure_rate_max_over_mean", 50.0 |
| ) |
| self.uniform_sampling_rate = self.adaptive_sampling_cfg.get("uniform_sampling_rate", 0.1) |
|
|
| |
| |
| self.max_prob_per_bin_cfg = self.adaptive_sampling_cfg.get("max_prob_per_bin", None) |
| self.max_prob_per_motion_cfg = self.adaptive_sampling_cfg.get("max_prob_per_motion", None) |
| self.adp_samp_num_failures = ( |
| torch.ones(self.adp_samp_num_bins, device=self._device, dtype=torch.float32) |
| * init_num_failures |
| ) |
| self.adp_samp_num_episodes = ( |
| torch.ones(self.adp_samp_num_bins, device=self._device, dtype=torch.float32) |
| * init_num_failures |
| ) |
| self.adp_samp_failure_rate = torch.ones( |
| self.adp_samp_num_bins, device=self._device, dtype=torch.float32 |
| ) |
| self.adp_samp_failure_rate_raw = torch.ones( |
| self.adp_samp_num_bins, device=self._device, dtype=torch.float32 |
| ) |
| self.adp_sampling_prob = ( |
| torch.ones(self.adp_samp_num_bins, device=self._device, dtype=torch.float64) |
| / self.adp_samp_num_bins |
| ) |
|
|
| def get_state_dict(self): |
| """Return a serializable state dict for checkpointing adaptive sampling stats. |
| |
| Returns: |
| Dict containing ``adp_samp_num_episodes`` and ``adp_samp_num_failures`` |
| tensors if adaptive sampling is enabled, otherwise an empty dict. |
| """ |
| state_dict = {} |
| if self.use_adaptive_sampling: |
| state_dict.update( |
| { |
| "adp_samp_num_episodes": self.adp_samp_num_episodes, |
| "adp_samp_num_failures": self.adp_samp_num_failures, |
| } |
| ) |
| return state_dict |
|
|
| def load_state_dict(self, state_dict): |
| """Restore adaptive sampling statistics from a checkpoint. |
| |
| Validates that the bin count matches before restoring. If it does not match |
| (e.g. dataset changed between runs), the load is silently skipped. |
| |
| Args: |
| state_dict: Dict previously returned by ``get_state_dict()``. |
| """ |
| if self.use_adaptive_sampling and "adp_samp_num_episodes" in state_dict: |
| if len(self.adp_samp_num_failures) != len(state_dict["adp_samp_num_failures"]): |
| print("Adaptive sampling state dict does not match. Skipping load.") |
| return |
|
|
| self.adp_samp_num_episodes[:] = state_dict["adp_samp_num_episodes"].to(self._device) |
| self.adp_samp_num_failures[:] = state_dict["adp_samp_num_failures"].to(self._device) |
| self.sync_and_compute_adaptive_sampling(sync_across_gpus=False) |
| return |
|
|
| def update_adaptive_sampling(self, failure, motion_ids, motion_time_steps): |
| """Update adaptive sampling statistics based on training outcomes. |
| |
| Increments episode counts for all sampled bins, and failure counts for bins |
| where the policy terminated early. Uses bincount for efficient batched updates |
| when multiple environments hit the same bin. |
| |
| Args: |
| failure: Boolean tensor of shape ``(N,)`` indicating which environments |
| terminated due to failure (not timeout). |
| motion_ids: Tensor of shape ``(N,)`` with batch-local motion indices. |
| motion_time_steps: Tensor of shape ``(N,)`` with the simulation time step |
| at which the episode ended (or was sampled). |
| """ |
| |
| dataset_motion_ids = self.get_motion_ids_in_dataset(motion_ids) |
|
|
| time_steps = self.adp_samp_length_starts[dataset_motion_ids] + motion_time_steps |
|
|
| |
| if len(time_steps) > 0: |
| |
| bin_ids = self.adp_samp_frame_to_bin[time_steps] |
| counts = torch.bincount(bin_ids, minlength=self.adp_samp_num_bins) |
| counts = counts / self.adp_samp_bin_motion_length |
| self.adp_samp_num_episodes += counts |
|
|
| |
| if failure.any(): |
| failed_time_steps = time_steps[failure] |
| |
| if len(failed_time_steps) > 0: |
| bin_ids = self.adp_samp_frame_to_bin[failed_time_steps] |
| failure_counts = torch.bincount(bin_ids, minlength=self.adp_samp_num_bins) |
| failure_counts_multiplier = self.adaptive_sampling_cfg.get( |
| "failure_counts_multiplier", 1 |
| ) |
| self.adp_samp_num_failures += failure_counts * failure_counts_multiplier |
|
|
| def sync_and_compute_adaptive_sampling(self, accelerator=None, sync_across_gpus=False): |
| """Synchronize adaptive sampling stats across GPUs and recompute probabilities. |
| |
| In multi-GPU training, averages episode/failure counts across all processes |
| before recomputing the per-bin sampling distribution. Optionally applies |
| failure-rate decay to propagate difficulty information to preceding bins. |
| |
| Args: |
| accelerator: HuggingFace Accelerator instance for multi-GPU gather. |
| Required when ``sync_across_gpus=True``. |
| sync_across_gpus: Whether to synchronize statistics across distributed |
| processes before computing probabilities. |
| """ |
| if not self.use_adaptive_sampling: |
| return |
|
|
| if sync_across_gpus: |
| with common.Timer("sync_adaptive_sampling_across_gpus"): |
| adp_samp_stats = torch.cat( |
| [self.adp_samp_num_episodes, self.adp_samp_num_failures], dim=-1 |
| ) |
| adp_samp_stats_all = accelerator.gather(adp_samp_stats).reshape( |
| -1, *adp_samp_stats.shape |
| ) |
| adp_samp_stats_all = adp_samp_stats_all.mean(dim=0) |
| self.adp_samp_num_episodes, self.adp_samp_num_failures = adp_samp_stats_all.chunk( |
| 2, dim=-1 |
| ) |
|
|
| with common.Timer("compute_sampling_prob"): |
| failure_rate = self.adp_samp_num_failures / self.adp_samp_num_episodes |
| self.adp_samp_failure_rate_raw = failure_rate.clone() |
| self.adp_samp_failure_rate = failure_rate |
| |
| |
| if self.adaptive_sampling_cfg.get("use_failure_rate_decay", False): |
| gamma = self.adaptive_sampling_cfg.get("decay_gamma", 0.99) |
| num_steps = self.adp_samp_num_episodes.shape[0] |
| failure_rate_w_decay = torch.zeros_like(failure_rate) |
| for step in reversed(range(num_steps)): |
| if step == num_steps - 1: |
| next_failure_rate = 0 |
| next_is_not_terminal = 0.0 |
| else: |
| next_failure_rate = failure_rate_w_decay[step + 1] |
| next_is_not_terminal = ( |
| 1.0 - self.adp_samp_bin_new_motion_mask[step + 1].float() |
| ) |
| failure_rate_w_decay[step] = ( |
| failure_rate[step] + next_is_not_terminal * gamma * next_failure_rate |
| ) |
| self.adp_samp_failure_rate = failure_rate_w_decay |
|
|
| |
| self.update_adaptive_sampling_probabilities() |
| return |
|
|
| def update_adaptive_sampling_probabilities(self): |
| """Recompute per-bin sampling probabilities for the currently loaded motion batch. |
| |
| Blends failure-rate-based probabilities with a uniform baseline (controlled by |
| ``uniform_sampling_rate``), then applies optional max-probability constraints |
| per bin and per motion to prevent over-concentration on outlier sequences. |
| See the inline comments for detailed rationale on the constraint design. |
| """ |
| self.adp_samp_failure_rate = self.adp_samp_failure_rate.double() |
| self.adp_samp_active_failure_rate = self.adp_samp_failure_rate[ |
| self.adp_samp_active_motion_bins |
| ] |
| adp_samp_failure_rate_upper_bound = ( |
| self.adp_samp_active_failure_rate.mean() * self.adp_samp_failure_rate_max_over_mean |
| ) |
| adp_samp_active_failure_rate_clipped = torch.clip( |
| self.adp_samp_active_failure_rate, 0.0, adp_samp_failure_rate_upper_bound |
| ) |
| failure_based_sampling_prob = ( |
| adp_samp_active_failure_rate_clipped / adp_samp_active_failure_rate_clipped.sum() |
| ) |
| uniform_sampling_prob = torch.ones_like(failure_based_sampling_prob) / len( |
| failure_based_sampling_prob |
| ) |
| self.adp_sampling_active_prob = ( |
| failure_based_sampling_prob * (1 - self.uniform_sampling_rate) |
| + uniform_sampling_prob * self.uniform_sampling_rate |
| ) |
| self.adp_sampling_active_prob *= self.adp_samp_bin_weights[self.adp_samp_active_motion_bins] |
| self.adp_sampling_active_prob = ( |
| self.adp_sampling_active_prob / self.adp_sampling_active_prob.sum() |
| ) |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| |
| if self.max_prob_per_bin_cfg is None and self.max_prob_per_motion_cfg is None: |
| self.adp_sampling_active_prob = self.adp_sampling_active_prob.float() |
| assert (self.adp_sampling_active_prob >= 0).all() |
| return |
|
|
| num_active_bins = len(self.adp_samp_active_motion_bins) |
| active_orig_motion_ids = self.adp_samp_bins[self.adp_samp_active_motion_bins, 0] |
| num_active_motions = len(active_orig_motion_ids.unique()) |
|
|
| |
| if self.max_prob_per_bin_cfg is not None: |
| if self.max_prob_per_bin_cfg == "auto": |
| |
| multiplier = self.adp_samp_failure_rate_max_over_mean |
| max_prob_per_bin = multiplier / num_active_bins if num_active_bins > 0 else 1.0 |
| else: |
| max_prob_per_bin = ( |
| float(self.max_prob_per_bin_cfg) if self.max_prob_per_bin_cfg else 0.0 |
| ) |
|
|
| |
| if max_prob_per_bin > 0 and num_active_bins > 1.0 / max_prob_per_bin: |
| self.adp_sampling_active_prob = torch.clamp( |
| self.adp_sampling_active_prob, max=max_prob_per_bin |
| ) |
| self.adp_sampling_active_prob = ( |
| self.adp_sampling_active_prob / self.adp_sampling_active_prob.sum() |
| ) |
|
|
| |
| if self.max_prob_per_motion_cfg is not None: |
| if self.max_prob_per_motion_cfg == "auto": |
| |
| multiplier = self.adp_samp_failure_rate_max_over_mean |
| max_prob_per_motion = ( |
| multiplier / num_active_motions if num_active_motions > 0 else 1.0 |
| ) |
| else: |
| max_prob_per_motion = ( |
| float(self.max_prob_per_motion_cfg) if self.max_prob_per_motion_cfg else 0.0 |
| ) |
|
|
| |
| if max_prob_per_motion > 0 and num_active_motions > 1.0 / max_prob_per_motion: |
| unique_motions = active_orig_motion_ids.unique() |
|
|
| for motion_id in unique_motions: |
| motion_mask = active_orig_motion_ids == motion_id |
| motion_total_prob = self.adp_sampling_active_prob[motion_mask].sum() |
|
|
| if motion_total_prob > max_prob_per_motion: |
| |
| scale_factor = max_prob_per_motion / motion_total_prob |
| self.adp_sampling_active_prob[motion_mask] *= scale_factor |
|
|
| |
| self.adp_sampling_active_prob = ( |
| self.adp_sampling_active_prob / self.adp_sampling_active_prob.sum() |
| ) |
| |
|
|
| self.adp_sampling_active_prob = self.adp_sampling_active_prob.float() |
| assert (self.adp_sampling_active_prob >= 0).all() |
|
|
| def update_adaptive_sampling_motion_sequences(self): |
| """Recompute global (full-dataset) motion-level sampling probabilities. |
| |
| Called before ``load_motions()`` to determine which motions to load next. |
| Aggregates per-bin failure rates into per-motion probabilities and applies |
| the same max-probability constraints as the batch-level update. |
| """ |
| self.adp_samp_failure_rate = self.adp_samp_failure_rate.double() |
| adp_samp_failure_rate_upper_bound = ( |
| self.adp_samp_failure_rate.mean() * self.adp_samp_failure_rate_max_over_mean |
| ) |
| adp_samp_failure_rate_clipped = torch.clip( |
| self.adp_samp_failure_rate, 0.0, adp_samp_failure_rate_upper_bound |
| ) |
| failure_based_sampling_prob = ( |
| adp_samp_failure_rate_clipped / adp_samp_failure_rate_clipped.sum() |
| ) |
| uniform_sampling_prob = torch.ones_like(failure_based_sampling_prob) / len( |
| failure_based_sampling_prob |
| ) |
| self.adp_sampling_prob = ( |
| failure_based_sampling_prob * (1 - self.uniform_sampling_rate) |
| + uniform_sampling_prob * self.uniform_sampling_rate |
| ) |
| self.adp_sampling_prob *= self.adp_samp_bin_weights |
| self.adp_sampling_prob = self.adp_sampling_prob / self.adp_sampling_prob.sum() |
|
|
| |
| |
| |
| |
| |
| if self.max_prob_per_bin_cfg is None and self.max_prob_per_motion_cfg is None: |
| |
| motion_sampling_probs = torch.zeros(self._num_unique_motions, device=self._device) |
| for orig_motion_id in range(self._num_unique_motions): |
| motion_sampling_probs[orig_motion_id] = self.adp_sampling_prob[ |
| self.orig_motion_id_to_bins[orig_motion_id] |
| ].sum() |
| self._sampling_prob = motion_sampling_probs / motion_sampling_probs.sum() |
| return |
|
|
| num_bins = self.adp_samp_num_bins |
| num_motions = self._num_unique_motions |
|
|
| |
| if self.max_prob_per_bin_cfg is not None: |
| if self.max_prob_per_bin_cfg == "auto": |
| |
| multiplier = self.adp_samp_failure_rate_max_over_mean |
| max_prob_per_bin = multiplier / num_bins if num_bins > 0 else 1.0 |
| else: |
| max_prob_per_bin = ( |
| float(self.max_prob_per_bin_cfg) if self.max_prob_per_bin_cfg else 0.0 |
| ) |
|
|
| |
| if max_prob_per_bin > 0 and num_bins > 1.0 / max_prob_per_bin: |
| self.adp_sampling_prob = torch.clamp(self.adp_sampling_prob, max=max_prob_per_bin) |
| self.adp_sampling_prob = self.adp_sampling_prob / self.adp_sampling_prob.sum() |
|
|
| |
| motion_sampling_probs = torch.zeros(self._num_unique_motions, device=self._device) |
| for orig_motion_id in range(self._num_unique_motions): |
| motion_sampling_probs[orig_motion_id] = self.adp_sampling_prob[ |
| self.orig_motion_id_to_bins[orig_motion_id] |
| ].sum() |
|
|
| |
| if self.max_prob_per_motion_cfg is not None: |
| if self.max_prob_per_motion_cfg == "auto": |
| |
| multiplier = self.adp_samp_failure_rate_max_over_mean |
| max_prob_per_motion = multiplier / num_motions if num_motions > 0 else 1.0 |
| else: |
| max_prob_per_motion = ( |
| float(self.max_prob_per_motion_cfg) if self.max_prob_per_motion_cfg else 0.0 |
| ) |
|
|
| |
| if max_prob_per_motion > 0 and num_motions > 1.0 / max_prob_per_motion: |
| motion_sampling_probs = torch.clamp(motion_sampling_probs, max=max_prob_per_motion) |
|
|
| self._sampling_prob = motion_sampling_probs / motion_sampling_probs.sum() |
|
|
| def update_adaptive_sampling_motion_frames(self): |
| """Build the active-bin index for the currently loaded motion batch. |
| |
| Maps each loaded motion to its corresponding global bins, creating |
| ``adp_samp_active_motion_bins`` which is used by |
| ``sample_motion_ids_and_time_steps()`` and |
| ``update_adaptive_sampling_probabilities()`` to sample and update only |
| the bins that correspond to currently loaded motions. |
| """ |
| self.adp_samp_active_motion_bins = [] |
| self.orig_motion_id_to_motion_ids = torch.zeros( |
| self._num_unique_motions, device=self._device, dtype=torch.long |
| ) |
| for motion_id, orig_motion_id in enumerate(self._curr_motion_ids): |
| bins = self.orig_motion_id_to_bins[orig_motion_id.item()] |
| self.adp_samp_active_motion_bins.append(bins) |
| self.orig_motion_id_to_motion_ids[orig_motion_id.item()] = motion_id |
|
|
| |
| adp_frames = self.adp_samp_num_frames[orig_motion_id].item() |
| loaded_frames = self._motion_num_frames[motion_id].item() |
| assert adp_frames == loaded_frames, ( |
| f"Adaptive sampling frame count mismatch for motion " |
| f"{orig_motion_id.item()} (key={self._motion_data_keys[orig_motion_id]}): " |
| f"adp_samp={adp_frames}, loaded={loaded_frames}. " |
| f"This means init_adaptive_sampling computed a different frame count " |
| f"than fk_batch produced at load time." |
| ) |
|
|
| self.adp_samp_active_motion_bins = torch.cat(self.adp_samp_active_motion_bins, dim=0) |
| self.update_adaptive_sampling_probabilities() |
|
|
| def sample_motion_ids_and_time_steps(self, n): |
| """Sample motion IDs and time steps using adaptive sampling probabilities. |
| |
| Draws bins from the active-bin distribution, then samples a random frame |
| within each selected bin. Optionally shifts the sampled frame backward by |
| a random offset (``pre_failure_sample_window``) so the policy starts |
| practicing before the difficult segment. |
| |
| Args: |
| n: Number of (motion_id, time_step) pairs to sample. |
| |
| Returns: |
| Tuple of (motion_ids, motion_time_steps) where: |
| - motion_ids: ``(n,)`` long tensor with batch-local motion indices. |
| - motion_time_steps: ``(n,)`` int tensor with frame indices. |
| """ |
| sampled_bin_ids = torch.multinomial( |
| self.adp_sampling_active_prob, num_samples=n, replacement=True |
| ).to(self._device) |
| bin_ids = self.adp_samp_active_motion_bins[sampled_bin_ids] |
| bins = self.adp_samp_bins[bin_ids] |
| orig_motion_ids, bin_start, bin_end = bins[:, 0], bins[:, 1], bins[:, 2] |
| motion_ids = self.orig_motion_id_to_motion_ids[orig_motion_ids] |
|
|
| motion_time_steps = ( |
| torch.rand(len(bin_start), device=bin_start.device) * (bin_end - bin_start) |
| ).floor().long() + bin_start |
| |
| pre_failure_sample_window = self.adaptive_sampling_cfg.get("pre_failure_sample_window", 0) |
| if pre_failure_sample_window > 0: |
| offset = torch.randint(pre_failure_sample_window, (n,), device=self._device) |
| motion_time_steps = (motion_time_steps - offset).clamp_min(0) |
| return motion_ids, motion_time_steps.int() |
|
|
