from __future__ import annotations import argparse import os import sys import time from pathlib import Path import numpy as np import torch # Optional external checkouts for legacy bundles. The released native H2 bundle # uses native_closd_inference and does not load these checkouts. MOTIONBRICKS_ROOT = Path( os.environ.get("WBT_MOTIONBRICKS_ROOT", "").strip() or str(Path(__file__).resolve().parents[2] / "motionbricks") ) CLOSD_ROOT = Path(os.environ.get("WBT_CLOSD_ROOT", "").strip() or str(Path(__file__).resolve().parents[2] / "CLoSD")) def _ensure_motionbricks_paths(root: str | Path = MOTIONBRICKS_ROOT) -> Path: root = Path(root).expanduser() for path in (root, root / "scripts"): path_s = str(path) if path_s not in sys.path: sys.path.insert(0, path_s) return root def _resolve_existing_relative(path: str | Path, *bases: Path) -> Path: path = Path(str(path)).expanduser() if path.is_absolute(): return path for base in bases: candidate = base / path if candidate.exists(): return candidate return bases[0] / path def _sync_cuda(device: torch.device) -> None: if torch.device(device).type == "cuda" and torch.cuda.is_available(): torch.cuda.synchronize(device) def _normalize_quat_wxyz_torch(quat: torch.Tensor) -> torch.Tensor: norm = torch.linalg.norm(quat, dim=-1, keepdim=True) safe = torch.clamp(norm, min=1.0e-8) out = quat / safe identity = torch.zeros_like(out) identity[..., 0] = 1.0 return torch.where(norm < 1.0e-6, identity, out) def _quat_mul_wxyz_torch(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: aw, ax, ay, az = a.unbind(dim=-1) bw, bx, by, bz = b.unbind(dim=-1) return torch.stack( [ aw * bw - ax * bx - ay * by - az * bz, aw * bx + ax * bw + ay * bz - az * by, aw * by - ax * bz + ay * bw + az * bx, aw * bz + ax * by - ay * bx + az * bw, ], dim=-1, ) def _align_qpos_batch_torch(qpos: torch.Tensor, source: torch.Tensor, target: torch.Tensor) -> torch.Tensor: out = qpos.clone() source_yaw = torch.atan2( 2.0 * (source[:, 3] * source[:, 6] + source[:, 4] * source[:, 5]), 1.0 - 2.0 * (source[:, 5] * source[:, 5] + source[:, 6] * source[:, 6]), ) target_yaw = torch.atan2( 2.0 * (target[:, 3] * target[:, 6] + target[:, 4] * target[:, 5]), 1.0 - 2.0 * (target[:, 5] * target[:, 5] + target[:, 6] * target[:, 6]), ) delta_yaw = target_yaw - source_yaw c, s = torch.cos(delta_yaw), torch.sin(delta_yaw) rel = out[:, :, :2] - source[:, None, :2] out[:, :, 0] = target[:, None, 0] + c[:, None] * rel[:, :, 0] - s[:, None] * rel[:, :, 1] out[:, :, 1] = target[:, None, 1] + s[:, None] * rel[:, :, 0] + c[:, None] * rel[:, :, 1] out[:, :, 2] += (target[:, 2] - source[:, 2])[:, None] delta_quat = torch.zeros((qpos.shape[0], 4), dtype=qpos.dtype, device=qpos.device) delta_quat[:, 0] = torch.cos(0.5 * delta_yaw) delta_quat[:, 3] = torch.sin(0.5 * delta_yaw) out[:, :, 3:7] = _normalize_quat_wxyz_torch(_quat_mul_wxyz_torch(delta_quat[:, None, :], out[:, :, 3:7])) return out def _translate_qpos_xy_batch_torch( qpos: torch.Tensor, source: torch.Tensor, target: torch.Tensor ) -> torch.Tensor: """Place a table-frame motion at target XY without changing yaw or height.""" out = qpos.clone() delta_xy = target[:, :2] - source[:, :2] out[:, :, :2] += delta_xy[:, None, :] return out def _finite_difference_torch(values: torch.Tensor, fps: float) -> torch.Tensor: vel = torch.zeros_like(values) if values.shape[1] > 1: vel[:, :-1] = (values[:, 1:] - values[:, :-1]) * float(fps) vel[:, -1] = vel[:, -2] return vel def _angular_velocity_wxyz_torch(quat_wxyz: torch.Tensor, fps: float) -> torch.Tensor: ang = torch.zeros((*quat_wxyz.shape[:3], 3), dtype=quat_wxyz.dtype, device=quat_wxyz.device) if quat_wxyz.shape[1] <= 1: return ang q0 = quat_wxyz[:, :-1] q1 = quat_wxyz[:, 1:] inv0 = q0.clone() inv0[..., 1:] *= -1.0 delta = _normalize_quat_wxyz_torch(_quat_mul_wxyz_torch(q1, inv0)) delta = torch.where(delta[..., 0:1] < 0.0, -delta, delta) vec = delta[..., 1:] vec_norm = torch.linalg.norm(vec, dim=-1, keepdim=True) angle = 2.0 * torch.atan2(vec_norm, torch.clamp(delta[..., 0:1], min=1.0e-8)) rotvec = torch.where(vec_norm > 1.0e-8, vec / torch.clamp(vec_norm, min=1.0e-8) * angle, 2.0 * vec) ang[:, :-1] = rotvec * float(fps) ang[:, -1] = ang[:, -2] return ang class MotionBricksCLoSDRuntime: """Small runtime wrapper around the MotionBricks CLoSD online demo.""" def __init__( self, model_path: str, *, motionbricks_root: str | Path = MOTIONBRICKS_ROOT, closd_root: str | Path = CLOSD_ROOT, hparams: str = "out/motionbricks_pose/version_1/hparams.yaml", skeleton_xml: str = "assets/skeletons/g1/g1.xml", humanoid_xml: str = "assets/skeletons/g1/g1_29dof_tt.xml", device: str = "cuda", device_id: int = 0, sampling_steps: int | None = None, fps: float = 50.0, condition_clamp_enabled: bool = False, condition_time_range: tuple[float, float] = (0.24, 0.54), condition_lateral_range: tuple[float, float] = (-0.80, 0.28), condition_forward_range: tuple[float, float] = (0.17, 0.77), condition_height_range: tuple[float, float] = (0.70, 1.30), robot_joint_names: list[str] | tuple[str, ...] | None = None, robot_default_joint_pos: torch.Tensor | np.ndarray | None = None, reference_body_names: list[str] | tuple[str, ...] | None = None, ): self.root = Path(motionbricks_root).expanduser() self.closd_root = Path(closd_root).expanduser() self.device = torch.device(device if device == "cpu" or torch.cuda.is_available() else "cpu") self.condition_clamp_enabled = bool(condition_clamp_enabled) self.condition_time_range = tuple(map(float, condition_time_range)) self.condition_lateral_range = tuple(map(float, condition_lateral_range)) self.condition_forward_range = tuple(map(float, condition_forward_range)) self.condition_height_range = tuple(map(float, condition_height_range)) model_path_p = Path(model_path).expanduser() from .native_closd_inference import is_native_runtime_bundle if is_native_runtime_bundle(model_path_p): from .native_closd_inference import ( fk_matrix_to_quaternion, load_native_runtime_bundle, motion_to_qpos, quaternion_to_matrix, ) self._matrix_to_quaternion = fk_matrix_to_quaternion self._quaternion_to_matrix = quaternion_to_matrix self._native_motion_to_qpos = motion_to_qpos self.train_args, self.dataset, self.model, self.diffusion = load_native_runtime_bundle( model_path_p, self.device ) self.sampling_steps = max(1, int(sampling_steps)) if sampling_steps is not None else None self.fps = float(getattr(self.dataset, "fps", fps)) self._is_robot_qpos = True self.pred_len = int(self.train_args.pred_len) self.context_len = int(self.train_args.context_len) self.condition_frame = str(getattr(self.train_args, "g1_condition_frame", "root")) self.compact_condition = bool(getattr(self.train_args, "g1_compact_condition", 0)) self._init_robot_qpos_runtime( robot_joint_names=robot_joint_names, robot_default_joint_pos=robot_default_joint_pos, reference_body_names=reference_body_names, ) return self.root = _ensure_motionbricks_paths(self.root) import mujoco from omegaconf import OmegaConf, open_dict from motionbricks.geometry.quaternions import matrix_to_quaternion from motionbricks.helper.mujoco_helper import forward_kinematics, get_mujoco_converter from motionbricks.helper.pl_util import load_motion_rep from motionbricks.motionlib.core.utils.rotations import quaternion_to_matrix from scripts.pingpong.closd.feature_subset import load_feature_subset_packer from scripts.pingpong.closd.visualize_closd_g1_pingpong_online import load_runtime from scripts.pingpong.reference.generate_reference_buffer import ( BODY_IDS_FOR_ISAACLAB_G1, MUJOCO_TO_ISAACLAB, ) self._forward_kinematics = forward_kinematics self._matrix_to_quaternion = matrix_to_quaternion self._quaternion_to_matrix = quaternion_to_matrix self.body_ids_for_isaaclab = np.asarray(BODY_IDS_FOR_ISAACLAB_G1, dtype=np.int32) self.mujoco_to_isaaclab = np.asarray(MUJOCO_TO_ISAACLAB, dtype=np.int64) self.isaaclab_to_mujoco = np.argsort(self.mujoco_to_isaaclab) args = argparse.Namespace( model_path=str(model_path_p), model_type="auto", cvae_random_z=0, sample_index=0, use_avg=1, device=str(self.device), device_id=int(device_id), closd_root=str(self.closd_root), hparams=hparams, skeleton_xml=skeleton_xml, humanoid_xml=humanoid_xml, fps=float(fps), ) self.train_args, self.dataset, self.model, self.diffusion = load_runtime(args, self.device) self.sampling_steps = max(1, int(sampling_steps)) if sampling_steps is not None else None if self.sampling_steps is not None and hasattr(self.diffusion, "sampling_steps"): self.diffusion.sampling_steps = self.sampling_steps self.fps = float(getattr(self.dataset, "fps", fps)) from closd.diffusion_planner.utils.robot_qpos import is_robot_qpos_dataset self._is_robot_qpos = bool(is_robot_qpos_dataset(self.dataset)) self.pred_len = int(self.train_args.pred_len) self.context_len = int(self.train_args.context_len) self.condition_frame = str(getattr(self.train_args, "g1_condition_frame", "root")) self.compact_condition = bool(getattr(self.train_args, "g1_compact_condition", 0)) if self._is_robot_qpos: self._init_robot_qpos_runtime( robot_joint_names=robot_joint_names, robot_default_joint_pos=robot_default_joint_pos, reference_body_names=reference_body_names, ) return hparams_path = self._resolve(hparams) hparams_cfg = OmegaConf.load(hparams_path) version_dir = hparams_path.parent with open_dict(hparams_cfg): hparams_cfg.fps = self.fps if "skeleton" in hparams_cfg and "folder" in hparams_cfg.skeleton: skeleton_dir = version_dir / "skeleton" hparams_cfg.skeleton.folder = str( skeleton_dir if skeleton_dir.exists() else _resolve_existing_relative(hparams_cfg.skeleton.folder, self.root, version_dir) ) stats_cfg = OmegaConf.select(hparams_cfg, "motion_rep.stats") if stats_cfg is not None and "folder" in stats_cfg: stats_dir = version_dir / "stats" / "motion" hparams_cfg.motion_rep.stats.folder = str( stats_dir if stats_dir.exists() else _resolve_existing_relative(stats_cfg.folder, self.root, version_dir) ) self.motion_rep = load_motion_rep(hparams_cfg).to(self.device) self.converter = get_mujoco_converter(self.motion_rep, str(self._resolve(skeleton_xml))).to(self.device) from whole_body_tracking.utils.external_deps import closd_dataset_remap # Checkpoints record the dataset location of the machine they were # trained on; remap it when WBT_CLOSD_DATASETS_ROOT points elsewhere. self.train_args.data_dir = closd_dataset_remap(self.train_args.data_dir) self.feature_packer = load_feature_subset_packer(self.train_args.data_dir, self.device) self.mj_model = mujoco.MjModel.from_xml_path(str(self._resolve(humanoid_xml))) mujoco_joint_names = [ mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_JOINT, joint_id) for joint_id in range(self.mj_model.njnt) if self.mj_model.jnt_type[joint_id] != mujoco.mjtJoint.mjJNT_FREE ] self.isaaclab_joint_names = np.asarray( [mujoco_joint_names[int(idx)] for idx in self.mujoco_to_isaaclab], dtype="U64", ) self.isaaclab_body_names = np.asarray( [ mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, int(body_id)) for body_id in self.body_ids_for_isaaclab ], dtype="U64", ) fk_body_names = [] for name in self.converter._mujoco_joint_including_root_list: base = str(name) if base.endswith("_skel"): base = base[: -len("_skel")] fk_body_names.append("pelvis" if base == "pelvis" else f"{base}_link") try: self.torch_fk_to_isaaclab_body = np.asarray( [fk_body_names.index(str(name)) for name in self.isaaclab_body_names.tolist()], dtype=np.int64, ) except ValueError: self.torch_fk_to_isaaclab_body = np.arange(len(self.isaaclab_body_names), dtype=np.int64) self.mujoco_to_isaaclab_t = torch.as_tensor(self.mujoco_to_isaaclab, dtype=torch.long, device=self.device) self.torch_fk_to_isaaclab_body_t = torch.as_tensor( self.torch_fk_to_isaaclab_body, dtype=torch.long, device=self.device ) self._fk_axes = self.converter._mujoco_joint_axis_values_mujoco_space.to( device=self.device, dtype=torch.float32 ) self._fk_motion_ids = self.converter._mujoco_indices_to_motion_indices.to( device=self.device, dtype=torch.long ) self._fk_rot_offsets = self.converter._rot_offsets_q2t.to(device=self.device, dtype=torch.float32)[ self._fk_motion_ids ] self._fk_neutral = self.converter._capture_neutral_joints_mujoco.to(device=self.device, dtype=torch.float32) self._fk_parents = self.converter._mujoco_joint_including_root_parent_list.to( device=self.device, dtype=torch.long ) self._fk_idx_levs = [ idx.to(device=self.device, dtype=torch.long) for idx in self.converter._mujoco_joint_idx_levs ] @staticmethod def _canonical_body_name(name: str) -> str: name = str(name).lower() return "pelvis" if name == "pelvis_link" else name def _init_robot_qpos_runtime( self, *, robot_joint_names: list[str] | tuple[str, ...] | None, robot_default_joint_pos: torch.Tensor | np.ndarray | None, reference_body_names: list[str] | tuple[str, ...] | None, ) -> None: """Initialize native qpos/FK metadata shared by G1/A3/H2 bundles.""" if robot_joint_names is None or reference_body_names is None: raise ValueError( "Robot-native CLoSD checkpoints require robot_joint_names and reference_body_names." ) model_joint_names = [str(name) for name in self.dataset.robot_dof_names] robot_joint_names = [str(name) for name in robot_joint_names] missing_joints = [name for name in model_joint_names if name not in robot_joint_names] if missing_joints: raise ValueError(f"CLoSD joints are missing from the IsaacLab robot: {missing_joints}") # This mapping selects the model's 29 joints from IsaacLab. H2 has two # additional actuated head joints; references for those stay at the # articulation defaults and are still exposed to the RL policy. self.isaaclab_to_mujoco = np.asarray( [robot_joint_names.index(name) for name in model_joint_names], dtype=np.int64 ) self._native_model_to_robot_t = torch.as_tensor( self.isaaclab_to_mujoco, dtype=torch.long, device=self.device ) self.isaaclab_joint_names = np.asarray(robot_joint_names, dtype="U64") if robot_default_joint_pos is None: robot_default_joint_pos = np.zeros(len(robot_joint_names), dtype=np.float32) default_joint_pos = torch.as_tensor( robot_default_joint_pos, dtype=torch.float32, device=self.device ).reshape(-1) if default_joint_pos.numel() != len(robot_joint_names): raise ValueError( f"Expected {len(robot_joint_names)} default joint positions, got {default_joint_pos.numel()}." ) self._native_default_joint_pos = default_joint_pos model_body_names = [str(name) for name in self.dataset.robot_body_names] canonical_to_index = { self._canonical_body_name(name): index for index, name in enumerate(model_body_names) } missing_bodies = [ name for name in reference_body_names if self._canonical_body_name(name) not in canonical_to_index ] if missing_bodies: raise ValueError(f"CLoSD FK bodies are missing from the dataset: {missing_bodies}") self.isaaclab_body_names = np.asarray([str(name) for name in reference_body_names], dtype="U64") self._native_reference_body_ids_t = torch.as_tensor( [canonical_to_index[self._canonical_body_name(name)] for name in reference_body_names], dtype=torch.long, device=self.device, ) self._native_parents = torch.as_tensor( self.dataset.robot_parent_indices, dtype=torch.long, device=self.device ) self._native_translations = torch.as_tensor( self.dataset.robot_local_translation, dtype=torch.float32, device=self.device ) if hasattr(self.dataset, "robot_fixed_rotation_matrix"): self._native_fixed_rotations = torch.as_tensor( self.dataset.robot_fixed_rotation_matrix, dtype=torch.float32, device=self.device, ) else: local_quat = torch.as_tensor( self.dataset.robot_local_rotation_wxyz, dtype=torch.float32, device=self.device, ) self._native_fixed_rotations = self._quaternion_to_matrix(local_quat) self._native_dof_indices = torch.as_tensor( self.dataset.robot_joint_dof_indices, dtype=torch.long, device=self.device ) self._native_joint_axes = torch.as_tensor( self.dataset.robot_joint_axes, dtype=torch.float32, device=self.device ) self.qpos_dim = 7 + len(model_joint_names) def _resolve(self, path: str | Path) -> Path: path = Path(path).expanduser() return path if path.is_absolute() else self.root / path def _qpos_to_context(self, qpos: torch.Tensor | np.ndarray) -> torch.Tensor: if torch.is_tensor(qpos): qpos_t = qpos.to(device=self.device, dtype=torch.float32) if qpos_t.ndim == 2: qpos_t = qpos_t.unsqueeze(0) else: qpos_np = np.asarray(qpos, dtype=np.float32) if qpos_np.ndim == 2: qpos_np = qpos_np[None] qpos_t = torch.as_tensor(qpos_np, dtype=torch.float32, device=self.device) if self._is_robot_qpos: root_rotation = self._quaternion_to_matrix(qpos_t[..., 3:7]) root_6d = torch.cat([root_rotation[..., :, 0], root_rotation[..., :, 1]], dim=-1) features = torch.cat([qpos_t[..., :3], root_6d, qpos_t[..., 7:]], dim=-1) mean = torch.as_tensor(self.dataset.motion_mean, dtype=features.dtype, device=self.device) std = torch.as_tensor(self.dataset.motion_std, dtype=features.dtype, device=self.device) features = (features - mean) / std.clamp_min(1.0e-8) return features.permute(0, 2, 1).unsqueeze(2).contiguous() with torch.no_grad(): joint_pos, joint_rot = self.converter.convert_mujoco_qpos_to_motion_transforms(qpos_t) lengths = torch.full((qpos_t.shape[0],), qpos_t.shape[1], dtype=torch.long, device=self.device) global_motion = self.motion_rep.dual_rep.global_motion_rep( {"posed_joints": joint_pos, "global_joint_rots": joint_rot}, to_normalize=True, lengths=lengths, ) global_motion = self.motion_rep.dual_rep.global_motion_rep.change_first_heading( global_motion, torch.zeros(qpos_t.shape[0], device=self.device), is_normalized=True, to_normalize=True, ) local_motion = self.motion_rep.dual_rep.global_to_local( global_motion, is_normalized=True, to_normalize=True, lengths=lengths, ) if self.feature_packer is not None and local_motion.shape[-1] == self.feature_packer.full_dim: local_motion = local_motion[..., self.feature_packer.subset_indices.to(local_motion.device)] return local_motion.permute(0, 2, 1).unsqueeze(2).contiguous() def _pad_context_qpos(self, qpos: np.ndarray) -> np.ndarray: qpos = np.asarray(qpos, dtype=np.float32) if qpos.shape[0] >= self.context_len: return qpos[-self.context_len :].copy() pad = np.repeat(qpos[:1], self.context_len - qpos.shape[0], axis=0) return np.concatenate([pad, qpos], axis=0).astype(np.float32, copy=False) def _pad_context_qpos_batch(self, qpos: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray: if torch.is_tensor(qpos): qpos_t = qpos.to(device=self.device, dtype=torch.float32) if qpos_t.ndim == 2: qpos_t = qpos_t.unsqueeze(0) if qpos_t.shape[1] >= self.context_len: return qpos_t[:, -self.context_len :].contiguous() pad = qpos_t[:, :1].expand(-1, self.context_len - qpos_t.shape[1], -1) return torch.cat([pad, qpos_t], dim=1).contiguous() qpos = np.asarray(qpos, dtype=np.float32) if qpos.ndim == 2: return self._pad_context_qpos(qpos)[None] if qpos.shape[1] >= self.context_len: return qpos[:, -self.context_len :].copy() pad = np.repeat(qpos[:, :1], self.context_len - qpos.shape[1], axis=1) return np.concatenate([pad, qpos], axis=1).astype(np.float32, copy=False) def _table_condition_batch( self, cmd: torch.Tensor, root_qpos: torch.Tensor | np.ndarray, default_time: float ) -> torch.Tensor: """Match G1PingPongDataset's fixed-table condition transform.""" cmd = cmd.reshape(-1, 9).to(device=self.device, dtype=torch.float32) root_qpos_t = torch.as_tensor(root_qpos, dtype=torch.float32, device=self.device) cond = cmd.clone() active = cond[:, 0] >= 0.5 cond[~active] = 0.0 cond[~active, 1] = float(default_time) if not bool(active.any()): return cond horizontal_delta = cond[active, 2:4] - root_qpos_t[active, :2] # MotionBricks table axes are [world_y, world_x, world_z]. Height is # deliberately absolute/ground-relative for the no-context model. cond[active, 2:5] = torch.stack( [horizontal_delta[:, 1], horizontal_delta[:, 0], cond[active, 4]], dim=-1 ) cond[active, 5:8] = cond[active, 5:8][:, [1, 0, 2]] return cond def _motion_to_qpos_batch_tensor(self, motion: torch.Tensor) -> torch.Tensor: if self._is_robot_qpos: with torch.no_grad(): if hasattr(self, "_native_motion_to_qpos"): return self._native_motion_to_qpos(motion, self.dataset).contiguous() from closd.diffusion_planner.utils.robot_qpos import motion_to_qpos return motion_to_qpos(motion, self.dataset).contiguous() features = motion.permute(0, 3, 1, 2).squeeze(-1).contiguous() if self.feature_packer is not None: features = self.feature_packer.pack(features) with torch.no_grad(): qpos = self.converter.convert_motion_features_to_mujoco_qpos( features, self.motion_rep.dual_rep.local_motion_rep, is_normalized=True, root_quat_w_first=True, ) return qpos.contiguous() def _world_to_local_cmd_batch( self, cmd: torch.Tensor, root_qpos: torch.Tensor | np.ndarray, default_time: float ) -> torch.Tensor: cmd = cmd.reshape(-1, 9).to(device=self.device, dtype=torch.float32) if torch.is_tensor(root_qpos): root_qpos_t = root_qpos.to(device=self.device, dtype=torch.float32) else: root_qpos_t = torch.as_tensor(root_qpos, dtype=torch.float32, device=self.device) cond = cmd.clone() active = cond[:, 0] >= 0.5 cond[~active] = 0.0 cond[~active, 1] = float(default_time) if not bool(active.any()): return cond root = root_qpos_t[active, :3] q = root_qpos_t[active, 3:7] yaw = torch.atan2( 2.0 * (q[:, 0] * q[:, 3] + q[:, 1] * q[:, 2]), 1.0 - 2.0 * (q[:, 2] * q[:, 2] + q[:, 3] * q[:, 3]), ) c = torch.cos(-yaw) s = torch.sin(-yaw) rel_pos = cond[active, 2:5] - root rel_vel = cond[active, 5:8] pos_x = c * rel_pos[:, 0] - s * rel_pos[:, 1] pos_y = s * rel_pos[:, 0] + c * rel_pos[:, 1] vel_x = c * rel_vel[:, 0] - s * rel_vel[:, 1] vel_y = s * rel_vel[:, 0] + c * rel_vel[:, 1] cond[active, 2:5] = torch.stack([pos_y, pos_x, rel_pos[:, 2]], dim=-1) cond[active, 5:8] = torch.stack([vel_y, vel_x, rel_vel[:, 2]], dim=-1) return cond def _clamp_local_condition(self, cond: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: raw_cond = cond.clone() if not self.condition_clamp_enabled: return raw_cond, cond active = cond[:, 0] >= 0.5 if not bool(active.any()): return raw_cond, cond cond = cond.clone() ranges = ( self.condition_time_range, self.condition_lateral_range, self.condition_forward_range, self.condition_height_range, ) for index, value_range in enumerate(ranges, start=1): lo, hi = sorted(value_range) cond[active, index] = torch.clamp(cond[active, index], min=lo, max=hi) return raw_cond, cond def _generate_batch_qpos_tensor( self, cmd_local: torch.Tensor | np.ndarray, *, prefix_qpos: torch.Tensor | np.ndarray, played_frames: np.ndarray | torch.Tensor | None = None, profile: bool = False, ) -> dict[str, object]: timings: dict[str, float] = {} total_t0 = time.perf_counter() if profile else 0.0 setup_t0 = time.perf_counter() if profile else 0.0 if torch.is_tensor(cmd_local): cmd = cmd_local.to(device=self.device, dtype=torch.float32).reshape(-1, 9) else: cmd = torch.as_tensor(cmd_local, dtype=torch.float32, device=self.device).reshape(-1, 9) batch_size = int(cmd.shape[0]) if batch_size == 0: return { "qpos": torch.zeros( (0, 0, self.qpos_dim if self._is_robot_qpos else self.mj_model.nq), dtype=torch.float32, device=self.device, ), "played_frames": torch.zeros((0,), dtype=torch.long, device=self.device), "_timing": timings, } prefix_qpos_t = torch.as_tensor(prefix_qpos, dtype=torch.float32, device=self.device) if prefix_qpos_t.ndim == 2: prefix_qpos_t = prefix_qpos_t.unsqueeze(0) if prefix_qpos_t.shape[1] < 1: raise ValueError("prefix_qpos must retain at least one placement/seed frame.") seed_qpos_t = prefix_qpos_t[:, -1] if self.context_len > 0: display_tail = self._pad_context_qpos_batch(prefix_qpos_t) context = self._qpos_to_context(display_tail) display_tail_t = display_tail.to(device=self.device, dtype=torch.float32) cond_root_qpos = display_tail_t[:, -1] else: # The model receives no prefix tokens, but the current robot pose is # still needed outside the network for table-relative XY and world # placement of the generated canonical motion. # Do not convert a single qpos frame merely to discover the feature # shape: MotionBricks' velocity feature path requires at least two # frames. A no-context model needs only a correctly shaped empty # prefix tensor. context = torch.empty( (batch_size, int(self.model.njoints), int(self.model.nfeats), 0), dtype=torch.float32, device=self.device, ) cond_root_qpos = seed_qpos_t if self.condition_frame == "table": cond = self._table_condition_batch( cmd, cond_root_qpos, float(self.train_args.pred_len) / float(self.fps) ) else: cond = self._world_to_local_cmd_batch( cmd, cond_root_qpos, float(self.train_args.pred_len) / float(self.fps) ) raw_cond, cond = self._clamp_local_condition(cond) # Compact checkpoints are trained on [time_to_hit, lateral, forward, z] # rather than the legacy 9-D [active, time, position, velocity, sign] # command. Keep the 9-D tensors above for clamp diagnostics, and only # project the tensor that is passed into the model. model_cond = cond[:, 1:5] if self.compact_condition else cond if profile: timings["setup"] = time.perf_counter() - setup_t0 _sync_cuda(self.device) model_t0 = time.perf_counter() y = { "mask": torch.ones((batch_size, 1, 1, self.pred_len), dtype=torch.bool, device=self.device), "lengths": torch.full((batch_size,), self.pred_len, dtype=torch.long, device=self.device), "prefix": context, "pingpong_cond": model_cond, } sample_kwargs = {} diffusion_steps = int(getattr(self.diffusion, "num_timesteps", 0)) if self.sampling_steps is not None and not hasattr(self.diffusion, "sampling_steps") and diffusion_steps > 0: sample_kwargs["skip_timesteps"] = max(0, diffusion_steps - self.sampling_steps) with torch.no_grad(): pred = self.diffusion.p_sample_loop( self.model, (batch_size, context.shape[1], 1, self.pred_len), clip_denoised=False, model_kwargs={"y": y}, progress=False, **sample_kwargs, ) if profile: _sync_cuda(self.device) timings["model"] = time.perf_counter() - model_t0 if pred.shape[-1] == self.context_len + self.pred_len: pred = pred[..., self.context_len :] if pred.shape[-1] != self.pred_len: raise RuntimeError(f"Expected generated suffix with {self.pred_len} frames, got {pred.shape[-1]}.") decode_t0 = time.perf_counter() if profile else 0.0 full = torch.cat([context, pred], dim=-1) raw_qpos = self._motion_to_qpos_batch_tensor(full) if profile: _sync_cuda(self.device) timings["decode"] = time.perf_counter() - decode_t0 align_t0 = time.perf_counter() if profile else 0.0 if self.context_len > 0: align_source = raw_qpos[:, self.context_len - 1] qpos_start = self.context_len else: align_source = raw_qpos[:, 0] qpos_start = 0 if self.context_len == 0 and self.condition_frame == "table": # The no-context table model owns absolute ground-relative Z and # fixed-table yaw. Only horizontal translation is required to # place its canonical motion at the current robot location. aligned_qpos = _translate_qpos_xy_batch_torch(raw_qpos, align_source, cond_root_qpos) else: aligned_qpos = _align_qpos_batch_torch(raw_qpos, align_source, cond_root_qpos) qpos = aligned_qpos[:, qpos_start:].contiguous() if played_frames is None: played_frames_t = torch.full((batch_size,), qpos.shape[1], dtype=torch.long, device=self.device) else: played_frames_t = torch.as_tensor(played_frames, dtype=torch.long, device=self.device).reshape(-1) if played_frames_t.numel() == 1 and batch_size != 1: played_frames_t = played_frames_t.expand(batch_size) if played_frames_t.numel() != batch_size: raise ValueError(f"played_frames must have {batch_size} values, got {played_frames_t.numel()}.") played_frames_t = played_frames_t.clamp(min=1, max=qpos.shape[1]) if profile: _sync_cuda(self.device) timings["align"] = time.perf_counter() - align_t0 timings["total_qpos"] = time.perf_counter() - total_t0 return { "qpos": qpos, "played_frames": played_frames_t, "raw_condition": raw_cond, "condition": cond, "_timing": timings, } def _qpos_batch_to_isaaclab_motion_tensors(self, qpos: torch.Tensor, fps: int) -> dict[str, torch.Tensor]: qpos = qpos.to(device=self.device, dtype=torch.float32) if qpos.ndim == 2: qpos = qpos.unsqueeze(0) if self._is_robot_qpos: return self._robot_qpos_batch_to_isaaclab_motion_tensors(qpos, fps) batch_size, num_frames = qpos.shape[:2] device, dtype = qpos.device, qpos.dtype joint_pos = qpos[:, :, 7 : 7 + len(self.mujoco_to_isaaclab)][:, :, self.mujoco_to_isaaclab_t].contiguous() joint_vel = _finite_difference_torch(joint_pos, float(fps)) root_translation_mujoco = qpos[:, :, :3] root_quat_mujoco = qpos[:, :, 3:7] root_rotation_mujoco = self._quaternion_to_matrix(root_quat_mujoco) dof = qpos[:, :, 7:] zeros = torch.zeros_like(dof) quaternion_if_x_axis = torch.stack([torch.cos(dof / 2), torch.sin(dof / 2), zeros, zeros], dim=-1) quaternion_if_y_axis = torch.stack([torch.cos(dof / 2), zeros, torch.sin(dof / 2), zeros], dim=-1) quaternion_if_z_axis = torch.stack([torch.cos(dof / 2), zeros, zeros, torch.sin(dof / 2)], dim=-1) quaternion_from_xyz_axis = torch.stack( [quaternion_if_x_axis, quaternion_if_y_axis, quaternion_if_z_axis], dim=-1, ) axes = self._fk_axes.to(dtype=dtype) joint_quaternion = (quaternion_from_xyz_axis * axes[None, None, :, None, :]).sum(dim=-1) joint_rotation_matrix = self._quaternion_to_matrix(joint_quaternion) joint_rotation_matrix = torch.matmul(self._fk_rot_offsets.to(dtype=dtype)[None, None], joint_rotation_matrix) rot_matrices = torch.cat([root_rotation_mujoco[:, :, None], joint_rotation_matrix], dim=2) rot_matrices = rot_matrices.reshape(batch_size * num_frames, self.converter._nb_joints_mujoco, 3, 3) neutral = self._fk_neutral.to(dtype=dtype).expand( batch_size * num_frames, -1, -1, ) body_pos, body_rot = self._forward_kinematics(rot_matrices, neutral, self._fk_parents, self._fk_idx_levs, 0) body_pos = body_pos.reshape(batch_size, num_frames, self.converter._nb_joints_mujoco, 3) body_rot = body_rot.reshape(batch_size, num_frames, self.converter._nb_joints_mujoco, 3, 3) body_pos = body_pos + root_translation_mujoco[:, :, None, :] body_quat = _normalize_quat_wxyz_torch(self._matrix_to_quaternion(body_rot)) body_pos = body_pos[:, :, self.torch_fk_to_isaaclab_body_t] body_quat = body_quat[:, :, self.torch_fk_to_isaaclab_body_t] return { "joint_pos": joint_pos, "joint_vel": joint_vel, "body_pos_w": body_pos.contiguous(), "body_quat_w": body_quat.contiguous(), "body_lin_vel_w": _finite_difference_torch(body_pos, float(fps)).contiguous(), "body_ang_vel_w": _angular_velocity_wxyz_torch(body_quat, float(fps)).contiguous(), } def _robot_qpos_batch_to_isaaclab_motion_tensors( self, qpos: torch.Tensor, fps: int ) -> dict[str, torch.Tensor]: """Run dataset-provided differentiable FK and expand joints to IsaacLab order.""" batch_size, num_frames = qpos.shape[:2] dtype, device = qpos.dtype, qpos.device model_dof = qpos[..., 7:] joint_pos = self._native_default_joint_pos.to(dtype=dtype).view(1, 1, -1).expand( batch_size, num_frames, -1 ).clone() joint_pos[..., self._native_model_to_robot_t] = model_dof joint_vel = _finite_difference_torch(joint_pos, float(fps)) root_pos = qpos[..., :3] root_rot = self._quaternion_to_matrix(qpos[..., 3:7]) positions = [root_pos] rotations = [root_rot] identity = torch.eye(3, dtype=dtype, device=device).view(1, 1, 3, 3) for body_index in range(1, int(self._native_parents.numel())): parent = int(self._native_parents[body_index].item()) parent_pos = positions[parent] parent_rot = rotations[parent] translation = self._native_translations[body_index].to(dtype=dtype).view(1, 1, 3, 1) body_pos = parent_pos + torch.matmul(parent_rot, translation).squeeze(-1) local_rot = self._native_fixed_rotations[body_index].to(dtype=dtype).view(1, 1, 3, 3) dof_index = int(self._native_dof_indices[body_index].item()) if dof_index >= 0: axis = self._native_joint_axes[body_index].to(dtype=dtype) axis = axis / torch.linalg.norm(axis).clamp_min(1.0e-8) x, y, z = axis.unbind() zero = torch.zeros((), dtype=dtype, device=device) skew = torch.stack([zero, -z, y, z, zero, -x, -y, x, zero]).reshape(3, 3) outer = axis[:, None] * axis[None, :] angle = model_dof[..., dof_index] cosine = torch.cos(angle)[..., None, None] sine = torch.sin(angle)[..., None, None] joint_rot = cosine * identity + (1.0 - cosine) * outer + sine * skew local_rot = torch.matmul(local_rot, joint_rot) positions.append(body_pos) rotations.append(torch.matmul(parent_rot, local_rot)) body_pos = torch.stack(positions, dim=2)[:, :, self._native_reference_body_ids_t] body_rot = torch.stack(rotations, dim=2)[:, :, self._native_reference_body_ids_t] body_quat = _normalize_quat_wxyz_torch(self._matrix_to_quaternion(body_rot)) return { "joint_pos": joint_pos.contiguous(), "joint_vel": joint_vel.contiguous(), "body_pos_w": body_pos.contiguous(), "body_quat_w": body_quat.contiguous(), "body_lin_vel_w": _finite_difference_torch(body_pos, float(fps)).contiguous(), "body_ang_vel_w": _angular_velocity_wxyz_torch(body_quat, float(fps)).contiguous(), } def generate_batch_tensors( self, cmd_local: torch.Tensor | np.ndarray, *, prefix_qpos: torch.Tensor | np.ndarray, played_frames: np.ndarray | torch.Tensor | None = None, profile: bool = False, ) -> dict[str, torch.Tensor | np.ndarray | dict[str, float]]: generated = self._generate_batch_qpos_tensor( cmd_local, prefix_qpos=prefix_qpos, played_frames=played_frames, profile=profile, ) timings = dict(generated["_timing"]) export_t0 = time.perf_counter() if profile else 0.0 payload = self._qpos_batch_to_isaaclab_motion_tensors(generated["qpos"], int(round(self.fps))) if profile: _sync_cuda(self.device) timings["torch_fk_export"] = time.perf_counter() - export_t0 timings["total"] = sum( timings.get(key, 0.0) for key in ("setup", "model", "decode", "align", "torch_fk_export") ) played_frames_t = generated["played_frames"].to(device=self.device, dtype=torch.long) payload.update( { "joint_names": self.isaaclab_joint_names, "body_names": self.isaaclab_body_names, "_played_frames": played_frames_t, "_raw_condition": generated["raw_condition"], "_condition": generated["condition"], "_timing": timings, } ) return payload