Blame for runtime/motionbricks_closd_runtime.py

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