Blame for runtime/native_closd_inference.py

"""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