"""Self-contained CLoSD inference for G1/A3/H2 checkpoints. This module intentionally contains only the network and DDPM sampling path used by the compact, no-context, robot-qpos checkpoints. It has no imports from the external CLoSD or MotionBricks repositories. """ from __future__ import annotations import math from argparse import Namespace from pathlib import Path from types import SimpleNamespace import numpy as np import torch from torch import nn from torch.nn import functional as F NATIVE_BUNDLE_FORMAT = "table_tennis.native_closd_runtime.v1" ROBOT_QPOS_REPRESENTATION = "qpos_root_xyz_rot6d_dof_v1" G1_MOTIONBRICKS_REPRESENTATION = "motionbricks_g1_local_413d_v1" class PositionalEncoding(nn.Module): def __init__(self, latent_dim: int, dropout: float, max_len: int): super().__init__() self.dropout = nn.Dropout(p=dropout) pe = torch.zeros(max_len, latent_dim) position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1) div_term = torch.exp( torch.arange(0, latent_dim, 2, dtype=torch.float32) * (-math.log(10000.0) / latent_dim) ) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) self.register_buffer("pe", pe.unsqueeze(0).transpose(0, 1)) def forward(self, value: torch.Tensor) -> torch.Tensor: return self.dropout(value + self.pe[: value.shape[0]]) class TimestepEmbedder(nn.Module): def __init__(self, latent_dim: int, sequence_pos_encoder: PositionalEncoding): super().__init__() self.sequence_pos_encoder = sequence_pos_encoder self.time_embed = nn.Sequential( nn.Linear(latent_dim, latent_dim), nn.SiLU(), nn.Linear(latent_dim, latent_dim), ) def forward(self, timesteps: torch.Tensor) -> torch.Tensor: return self.time_embed(self.sequence_pos_encoder.pe[timesteps]).permute(1, 0, 2) class InputProcess(nn.Module): def __init__(self, input_features: int, latent_dim: int): super().__init__() self.poseEmbedding = nn.Linear(input_features, latent_dim) def forward(self, value: torch.Tensor) -> torch.Tensor: batch_size, num_joints, num_features, num_frames = value.shape value = value.permute(3, 0, 1, 2).reshape( num_frames, batch_size, num_joints * num_features ) return self.poseEmbedding(value) class OutputProcess(nn.Module): def __init__(self, latent_dim: int, num_joints: int, num_features: int): super().__init__() self.num_joints = num_joints self.num_features = num_features self.poseFinal = nn.Linear(latent_dim, num_joints * num_features) def forward(self, value: torch.Tensor) -> torch.Tensor: num_frames, batch_size, _ = value.shape value = self.poseFinal(value) value = value.reshape(num_frames, batch_size, self.num_joints, self.num_features) return value.permute(1, 2, 3, 0) class NativePingPongMDM(nn.Module): """The exact MDM subset used by A3/H2 table-condition checkpoints.""" def __init__(self, args: Namespace, motion_dim: int, condition_dim: int): super().__init__() if str(args.arch) != "trans_dec": raise ValueError(f"Native runtime only supports trans_dec, got {args.arch!r}.") if int(args.context_len) != 0: raise ValueError("Native runtime only supports no-context checkpoints.") self.njoints = int(motion_dim) self.nfeats = 1 self.pred_len = int(args.pred_len) self.context_len = int(args.context_len) self.mask_frames = bool(args.mask_frames) latent_dim = int(args.latent_dim) self.pingpong_cond_token = nn.Parameter(torch.zeros(1, 1, latent_dim)) self.input_process = InputProcess(self.njoints, latent_dim) self.sequence_pos_encoder = PositionalEncoding( latent_dim, float(getattr(args, "dropout", 0.1)), int(args.pos_embed_max_len) ) self.embed_pingpong_cond = nn.Sequential( nn.Linear(int(condition_dim), latent_dim), nn.SiLU(), nn.Linear(latent_dim, latent_dim), ) decoder_layer = nn.TransformerDecoderLayer( d_model=latent_dim, nhead=int(getattr(args, "num_heads", 4)), dim_feedforward=int(getattr(args, "ff_size", 1024)), dropout=float(getattr(args, "dropout", 0.1)), activation=str(getattr(args, "activation", "gelu")), ) self.seqTransDecoder = nn.TransformerDecoder(decoder_layer, num_layers=int(args.layers)) self.embed_timestep = TimestepEmbedder(latent_dim, self.sequence_pos_encoder) self.output_process = OutputProcess(latent_dim, self.njoints, self.nfeats) def forward( self, value: torch.Tensor, timesteps: torch.Tensor, y: dict[str, torch.Tensor] | None = None ) -> torch.Tensor: if y is None: raise ValueError("Native CLoSD inference requires a condition dictionary.") batch_size = value.shape[0] time_embedding = self.embed_timestep(timesteps) condition_embedding = self.embed_pingpong_cond(y["pingpong_cond"]).unsqueeze(0) # The exported A3/H2 checkpoints have context_len=0. Keep the empty # prefix concatenation so tensor layout matches the original model. value = torch.cat([y["prefix"], value], dim=-1) value = self.input_process(value) frame_mask = None if self.mask_frames and y["mask"].shape[-1] > 1: frame_mask = torch.logical_not( y["mask"][..., : value.shape[0]].squeeze(1).squeeze(1) ).to(device=value.device) value = self.sequence_pos_encoder(value) memory = torch.cat([time_embedding, condition_embedding], dim=0) output = self.seqTransDecoder( tgt=value, memory=memory, tgt_key_padding_mask=frame_mask, ) return self.output_process(output) def _cosine_betas(num_steps: int) -> np.ndarray: def alpha_bar(time_value: float) -> float: return math.cos((time_value + 0.008) / 1.008 * math.pi / 2) ** 2 betas = [] for index in range(num_steps): start = index / num_steps end = (index + 1) / num_steps betas.append(min(1.0 - alpha_bar(end) / alpha_bar(start), 0.999)) return np.asarray(betas, dtype=np.float64) def _extract(values: np.ndarray, timesteps: torch.Tensor, shape: tuple[int, ...]) -> torch.Tensor: result = torch.from_numpy(values).to(device=timesteps.device)[timesteps].float() while result.ndim < len(shape): result = result[..., None] return result.expand(shape) class NativeGaussianDiffusion: """Minimal START_X/FIXED_SMALL DDPM sampler used by the native bundles.""" def __init__(self, num_steps: int, noise_schedule: str, sigma_small: bool): if noise_schedule != "cosine": raise ValueError(f"Native runtime only supports cosine noise, got {noise_schedule!r}.") if not sigma_small: raise ValueError("Native runtime only supports sigma_small checkpoints.") self.num_timesteps = int(num_steps) betas = _cosine_betas(self.num_timesteps) alphas = 1.0 - betas alphas_cumprod = np.cumprod(alphas, axis=0) alphas_cumprod_prev = np.append(1.0, alphas_cumprod[:-1]) self.posterior_variance = ( betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod) ) self.posterior_log_variance_clipped = np.log( np.append(self.posterior_variance[1], self.posterior_variance[1:]) ) self.posterior_mean_coef1 = ( betas * np.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod) ) self.posterior_mean_coef2 = ( (1.0 - alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - alphas_cumprod) ) def p_sample_loop( self, model: nn.Module, shape: tuple[int, ...], *, clip_denoised: bool = False, model_kwargs: dict[str, dict[str, torch.Tensor]], progress: bool = False, skip_timesteps: int = 0, **_: object, ) -> torch.Tensor: if progress: raise ValueError("Native CLoSD runtime does not expose a progress bar.") device = next(model.parameters()).device value = torch.randn(*shape, device=device) batch_size = shape[0] for index in range(self.num_timesteps - int(skip_timesteps) - 1, -1, -1): timesteps = torch.full((batch_size,), index, dtype=torch.long, device=device) predicted_start = model(value, timesteps, **model_kwargs) if clip_denoised: predicted_start = predicted_start.clamp(-1.0, 1.0) mean = ( _extract(self.posterior_mean_coef1, timesteps, value.shape) * predicted_start + _extract(self.posterior_mean_coef2, timesteps, value.shape) * value ) log_variance = _extract( self.posterior_log_variance_clipped, timesteps, value.shape ) noise = torch.randn_like(value) nonzero = (timesteps != 0).float().view(batch_size, 1, 1, 1) value = mean + nonzero * torch.exp(0.5 * log_variance) * noise return torch.cat([model_kwargs["y"]["prefix"], value], dim=-1) def quaternion_to_matrix(quaternion: torch.Tensor) -> torch.Tensor: quaternion = quaternion / torch.linalg.norm(quaternion, dim=-1, keepdim=True).clamp_min(1.0e-8) real, i, j, k = quaternion.unbind(-1) two = 2.0 values = torch.stack( ( 1 - two * (j * j + k * k), two * (i * j - k * real), two * (i * k + j * real), two * (i * j + k * real), 1 - two * (i * i + k * k), two * (j * k - i * real), two * (i * k - j * real), two * (j * k + i * real), 1 - two * (i * i + j * j), ), dim=-1, ) return values.reshape(*quaternion.shape[:-1], 3, 3) def _copy_sign(magnitude: torch.Tensor, sign: torch.Tensor) -> torch.Tensor: return torch.where((magnitude < 0) != (sign < 0), -magnitude, magnitude) def _sqrt_positive(value: torch.Tensor) -> torch.Tensor: result = torch.zeros_like(value) positive = value > 0 result[positive] = torch.sqrt(value[positive]) return result def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor: m00, m11, m22 = matrix[..., 0, 0], matrix[..., 1, 1], matrix[..., 2, 2] real = 0.5 * _sqrt_positive(1 + m00 + m11 + m22) i = 0.5 * _sqrt_positive(1 + m00 - m11 - m22) j = 0.5 * _sqrt_positive(1 - m00 + m11 - m22) k = 0.5 * _sqrt_positive(1 - m00 - m11 + m22) return torch.stack( ( real, _copy_sign(i, matrix[..., 2, 1] - matrix[..., 1, 2]), _copy_sign(j, matrix[..., 0, 2] - matrix[..., 2, 0]), _copy_sign(k, matrix[..., 1, 0] - matrix[..., 0, 1]), ), dim=-1, ) def fk_matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor: """Match MotionBricks' candidate-selection convention for FK outputs.""" batch_shape = matrix.shape[:-2] m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( matrix.reshape(*batch_shape, 9), dim=-1 ) quaternion_abs = torch.sqrt( torch.stack( ( 1.0 + m00 + m11 + m22, 1.0 + m00 - m11 - m22, 1.0 - m00 + m11 - m22, 1.0 - m00 - m11 + m22, ), dim=-1, ).clamp_min(0.0) ) candidates = torch.stack( ( torch.stack( (quaternion_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01), dim=-1 ), torch.stack( (m21 - m12, quaternion_abs[..., 1] ** 2, m10 + m01, m02 + m20), dim=-1 ), torch.stack( (m02 - m20, m10 + m01, quaternion_abs[..., 2] ** 2, m12 + m21), dim=-1 ), torch.stack( (m10 - m01, m20 + m02, m21 + m12, quaternion_abs[..., 3] ** 2), dim=-1 ), ), dim=-2, ) floor = torch.tensor(0.1, dtype=matrix.dtype, device=matrix.device) candidates = candidates / (2.0 * quaternion_abs[..., None].maximum(floor)) selection = F.one_hot(quaternion_abs.argmax(dim=-1), num_classes=4)[..., None] return (selection * candidates).sum(dim=-2).reshape(*batch_shape, 4) def _cont6d_to_matrix(value: torch.Tensor) -> torch.Tensor: x_raw, y_raw = value[..., :3], value[..., 3:6] x = x_raw / torch.linalg.norm(x_raw, dim=-1, keepdim=True).clamp_min(1.0e-8) z = torch.cross(x, y_raw, dim=-1) z = z / torch.linalg.norm(z, dim=-1, keepdim=True).clamp_min(1.0e-8) y = torch.cross(z, x, dim=-1) return torch.stack([x, y, z], dim=-1) def _normalized_motion(model_output: torch.Tensor, dataset: SimpleNamespace) -> torch.Tensor: motion = model_output.squeeze(2).permute(0, 2, 1) mean = torch.as_tensor(dataset.motion_mean, dtype=motion.dtype, device=motion.device) std = torch.as_tensor(dataset.motion_std, dtype=motion.dtype, device=motion.device) epsilon = float(getattr(dataset, "motion_normalization_epsilon", 0.0)) scale = torch.sqrt(std.square() + epsilon) if epsilon > 0.0 else std return motion * scale + mean def _robot_qpos_motion_to_qpos( model_output: torch.Tensor, dataset: SimpleNamespace ) -> torch.Tensor: motion = _normalized_motion(model_output, dataset) root_quaternion = matrix_to_quaternion(_cont6d_to_matrix(motion[..., 3:9])) return torch.cat([motion[..., :3], root_quaternion, motion[..., 9:]], dim=-1) def _g1_motionbricks_motion_to_qpos( model_output: torch.Tensor, dataset: SimpleNamespace ) -> torch.Tensor: """Decode normalized MotionBricks local 413D features to G1 MuJoCo qpos.""" motion = _normalized_motion(model_output, dataset) device, dtype = motion.device, motion.dtype batch_size, num_frames = motion.shape[:2] # MotionBricks local-root layout begins with [yaw velocity, vx, vz, height]. # This G1 representation was trained with removing_heading=False, so vx/vz # are integrated directly in the motion frame before the fixed axis change. root_position = torch.zeros((batch_size, num_frames, 3), dtype=dtype, device=device) if num_frames > 1: root_position[:, 1:, 0] = motion[:, :-1, 1] / float(dataset.fps) root_position[:, 1:, 2] = motion[:, :-1, 2] / float(dataset.fps) root_position = torch.cumsum(root_position, dim=1) root_position[:, :, 1] = motion[:, :, 3] transform = torch.as_tensor( dataset.g1_motion_to_mujoco_matrix, dtype=dtype, device=device ) root_translation = torch.matmul( transform.view(1, 1, 3, 3), root_position[..., None] ).squeeze(-1) rotation_indices = torch.as_tensor( dataset.g1_global_rot_data_indices, dtype=torch.long, device=device ) rotation_6d = motion.index_select(-1, rotation_indices) joint_count = rotation_6d.shape[-1] // 6 global_rotation = _cont6d_to_matrix( rotation_6d.reshape(batch_size, num_frames, joint_count, 6) ) mujoco_to_motion_list = [int(value) for value in dataset.g1_mujoco_indices_to_motion_indices] parent_list = [int(value) for value in dataset.g1_mujoco_parent_indices] motion_parents = torch.arange(joint_count, dtype=torch.long, device=device) for child in range(1, len(parent_list)): child_motion = mujoco_to_motion_list[child - 1] parent = parent_list[child] motion_parents[child_motion] = 0 if parent == 0 else mujoco_to_motion_list[parent - 1] motion_parents[0] = 0 parent_rotation = global_rotation.index_select(2, motion_parents).clone() parent_rotation[:, :, 0] = torch.eye(3, dtype=dtype, device=device) local_rotation = parent_rotation.transpose(-1, -2) @ global_rotation offset_f2q = torch.as_tensor(dataset.g1_rot_offsets_f2q, dtype=dtype, device=device) local_rotation = offset_f2q @ local_rotation root_rotation = transform @ local_rotation[:, :, 0] @ transform.T root_quaternion = fk_matrix_to_quaternion(root_rotation) mujoco_to_motion = torch.as_tensor( mujoco_to_motion_list, dtype=torch.long, device=device ) joint_rotation = local_rotation.index_select(2, mujoco_to_motion) xyz_dof = torch.stack( [ torch.atan2(joint_rotation[..., 2, 1], joint_rotation[..., 2, 2]), torch.atan2(joint_rotation[..., 0, 2], joint_rotation[..., 0, 0]), torch.atan2(joint_rotation[..., 1, 0], joint_rotation[..., 1, 1]), ], dim=-1, ) axis_motion = torch.as_tensor(dataset.g1_joint_axis_motion, dtype=dtype, device=device) joint_dof = (xyz_dof * axis_motion.view(1, 1, -1, 3)).sum(dim=-1) return torch.cat([root_translation, root_quaternion, joint_dof], dim=-1) def motion_to_qpos(model_output: torch.Tensor, dataset: SimpleNamespace) -> torch.Tensor: representation = str(dataset.robot_representation) if representation == ROBOT_QPOS_REPRESENTATION: return _robot_qpos_motion_to_qpos(model_output, dataset) if representation == G1_MOTIONBRICKS_REPRESENTATION: return _g1_motionbricks_motion_to_qpos(model_output, dataset) raise ValueError(f"Unsupported native motion representation: {representation!r}") def is_native_runtime_bundle(path: str | Path) -> bool: return str(path).endswith(".closd-runtime.pt") def load_native_runtime_bundle( path: str | Path, device: torch.device ) -> tuple[Namespace, SimpleNamespace, NativePingPongMDM, NativeGaussianDiffusion]: bundle = torch.load(Path(path).expanduser(), map_location="cpu", weights_only=True) if bundle.get("format") != NATIVE_BUNDLE_FORMAT: raise ValueError(f"Unsupported native CLoSD bundle: {path}") args = Namespace(**bundle["args"]) dataset = SimpleNamespace(**bundle["dataset"]) if dataset.robot_representation not in ( ROBOT_QPOS_REPRESENTATION, G1_MOTIONBRICKS_REPRESENTATION, ): raise ValueError(f"Unsupported robot representation: {dataset.robot_representation!r}") condition_dim = 4 if bool(args.g1_compact_condition) else 9 model = NativePingPongMDM(args, int(dataset.motion_dim), condition_dim) missing, unexpected = model.load_state_dict(bundle["model_avg"], strict=False) if missing or unexpected: raise RuntimeError( f"Native CLoSD state mismatch: missing={list(missing)}, unexpected={list(unexpected)}" ) model.to(device).eval() diffusion = NativeGaussianDiffusion( int(args.diffusion_steps), str(args.noise_schedule), bool(args.sigma_small) ) return args, dataset, model, diffusion