kimodo/skeleton/base.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""Base skeleton class: hierarchy, joint metadata, and helpers for kinematics and motion."""

from pathlib import Path
from typing import Optional

import torch

from kimodo.assets import skeleton_asset_path

from .kinematics import fk
from .transforms import (
    from_standard_tpose,
    global_rots_to_local_rots,
    to_standard_tpose,
)


class SkeletonBase(torch.nn.Module):
    """Base class that stores a skeleton hierarchy and helper metadata.

    Subclasses define the static joint layout (joint names and parent links) and semantic groups
    (feet, hands, hips). This class builds index mappings, parent tensors, and convenience helpers
    used by kinematics, constraints, and motion conversion utilities.
    """

    # these should be defined in the subclass
    name = None
    bone_order_names_with_parents = None
    bone_order_names_no_root = None
    root_idx = None
    foot_joint_names = None
    foot_joint_idx = None
    hip_joint_names = None  # in order [right, left]
    hip_joint_idx = None  # in order [right, left]

    def __init__(
        self,
        folder: Optional[str] = None,
        name: Optional[str] = None,
        load: bool = True,
        **kwargs,  # to catch addition args in configs
    ):
        """Initialize a skeleton instance and optional neutral-pose assets.

        Args:
            folder: Folder containing serialized skeleton assets (for example
                `joints.p` and optional `standard_t_pose_global_offsets_rots.p`).
            name: Optional runtime name used to validate subclass compatibility.
            load: Whether to load tensor assets from `folder`.
            **kwargs: Unused extra config keys kept for config compatibility.
        """
        super().__init__()

        if name is not None:
            # Check that the name is not too far from the actual skeleton class name
            assert self.name in name
            self.name = name

        if folder is None:
            # Take the skeleton asset folder of the repo from the name
            # in case we don't override it
            folder = str(skeleton_asset_path(self.name))

        self.folder = folder

        self.dim = len(self.bone_order_names_with_parents)

        if load and folder is not None:
            pfolder = Path(folder)
            class_default_name = getattr(self.__class__, "name", None)
            print(
                "[kimodo][skeleton][init][entry]"
                f" class={self.__class__.__name__} name={self.name} class_default={class_default_name}"
                f" folder={pfolder} joints_exists={(pfolder / 'joints.p').exists()}"
            )
            if not (pfolder / "joints.p").exists():
                candidate_names = []
                if class_default_name:
                    candidate_names.append(str(class_default_name))
                if self.name:
                    candidate_names.append(str(self.name))
                # Robust fallback for renamed runtime names in model configs
                candidate_names.extend(["somaskel30", "somaskel77", "g1skel34", "smplx22"])
                for candidate in candidate_names:
                    fallback_folder = skeleton_asset_path(candidate)
                    if (fallback_folder / "joints.p").exists():
                        pfolder = fallback_folder
                        self.folder = str(pfolder)
                        print(
                            "[kimodo][skeleton][init][fallback]"
                            f" class={self.__class__.__name__} candidate={candidate} path={pfolder}"
                        )
                        break
            try:
                neutral_joints = torch.load(pfolder / "joints.p").squeeze()
            except Exception as error:
                print(
                    "[kimodo][skeleton][init][error]"
                    f" class={self.__class__.__name__} resolved_folder={pfolder}"
                    f" dir_exists={pfolder.exists()}"
                    f" dir_entries={sorted([p.name for p in pfolder.iterdir()]) if pfolder.exists() else []}"
                    f" error={type(error).__name__}: {error}"
                )
                raise
            self.register_buffer("neutral_joints", neutral_joints, persistent=False)

            if (pfolder / "bvh_joints.p").exists():
                bvh_neutral_joints = torch.load(pfolder / "bvh_joints.p").squeeze()
                self.register_buffer("bvh_neutral_joints", bvh_neutral_joints, persistent=False)

            global_offset_path = pfolder / "standard_t_pose_global_offsets_rots.p"
            if global_offset_path.exists():
                global_rot_offsets = torch.load(global_offset_path).squeeze()
                self.register_buffer("global_rot_offsets", global_rot_offsets, persistent=False)
            # Usefull for g1, where the rest pose is not zero
            baked_rest_path = pfolder / "rest_pose_local_rot.p"
            if baked_rest_path.exists():
                rest_pose_local_rot = torch.load(baked_rest_path).squeeze()
                self.register_buffer("rest_pose_local_rot", rest_pose_local_rot, persistent=False)
            print(
                "[kimodo][skeleton][init][exit]"
                f" class={self.__class__.__name__} resolved_folder={pfolder}"
                f" neutral_shape={tuple(self.neutral_joints.shape)}"
            )

        self.bone_order_names = [x for x, y in self.bone_order_names_with_parents]

        self.bone_parents = dict(self.bone_order_names_with_parents)
        self.bone_index = {x: idx for idx, x in enumerate(self.bone_order_names)}
        self.bone_order_names_index = self.bone_index

        # create the parents tensor on the fly
        joint_parents = torch.tensor(
            [-1 if (y := self.bone_parents[x]) is None else self.bone_index[y] for x in self.bone_order_names]
        )
        self.register_buffer("joint_parents", joint_parents, persistent=False)

        self.nbjoints = len(self.bone_order_names)

        # check lengths
        assert self.nbjoints == len(self.joint_parents)
        if "neutral_joints" in self.__dict__:
            assert self.nbjoints == len(self.neutral_joints)

        root_indices = torch.where(joint_parents == -1)[0]
        assert len(root_indices) == 1  # should be one root only
        self.root_idx = root_indices[0].item()

        if "neutral_joints" in self.__dict__:
            assert (self.neutral_joints[0] == 0).all()

        # remove the root
        self.bone_order_names_no_root = (
            self.bone_order_names[: self.root_idx] + self.bone_order_names[self.root_idx + 1 :]
        )

        self.foot_joint_names = self.left_foot_joint_names + self.right_foot_joint_names
        self.foot_joint_names_index = {x: idx for idx, x in enumerate(self.foot_joint_names)}

        self.left_foot_joint_idx = [
            self.bone_order_names.index(foot_joint) for foot_joint in self.left_foot_joint_names
        ]

        self.right_foot_joint_idx = [
            self.bone_order_names.index(foot_joint) for foot_joint in self.right_foot_joint_names
        ]

        self.foot_joint_idx = self.left_foot_joint_idx + self.right_foot_joint_idx

        self.hip_joint_idx = [self.bone_order_names.index(hip_joint) for hip_joint in self.hip_joint_names]

    def expand_joint_names(self, joint_names):
        """Expand base EE names [LeftFoot, RightFoot, LeftHand, RightHand] actual joint names to
        constrain position and rotations.

        Args:
            joint_names: list of list of base EE names to constrain

        Returns:
            rot_joint_names: list of list of joint names to constrain rotations
            pos_joint_names: list of list of joint names to constrain positions
        """

        base_ee = ["LeftFoot", "RightFoot", "LeftHand", "RightHand", "Hips"]

        pelvis_name = self.bone_order_names[self.root_idx]

        base_pos_names = [
            self.left_foot_joint_names,
            self.right_foot_joint_names,
            self.left_hand_joint_names,
            self.right_hand_joint_names,
            [pelvis_name],
        ]
        # base of each chain
        base_rot_names = [
            self.left_foot_joint_names[:1],
            self.right_foot_joint_names[:1],
            self.left_hand_joint_names[:1],
            self.right_hand_joint_names[:1],
            [pelvis_name],
        ]
        rot_joint_names = []
        pos_joint_names = []
        # loop through each EE joint group to constrain in the current keyframe
        for jname in joint_names:
            idx = base_ee.index(jname)
            rot_joint_names += base_rot_names[idx]
            pos_joint_names += base_pos_names[idx]
        return rot_joint_names, pos_joint_names

    def expand_joint_names_batched(self, joint_names):
        """Expand base EE names [LeftFoot, RightFoot, LeftHand, RightHand] actual joint names to
        constrain position and rotations.

        Args:
            joint_names: list of list of base EE names to constrain

        Returns:
            rot_joint_names: list of list of joint names to constrain rotations
            pos_joint_names: list of list of joint names to constrain positions
        """

        base_ee = ["LeftFoot", "RightFoot", "LeftHand", "RightHand", "Hips"]

        pelvis_name = self.bone_order_names[self.root_idx]

        base_pos_names = [
            self.left_foot_joint_names,
            self.right_foot_joint_names,
            self.left_hand_joint_names,
            self.right_hand_joint_names,
            [pelvis_name],
        ]
        # base of each chain
        base_rot_names = [
            self.left_foot_joint_names[:1],
            self.right_foot_joint_names[:1],
            self.left_hand_joint_names[:1],
            self.right_hand_joint_names[:1],
            [pelvis_name],
        ]
        # loop through each keyframe
        rot_joint_names = []
        pos_joint_names = []
        for key_joint_names in joint_names:
            key_rot_names = []
            key_pos_names = []
            # loop through each EE joint group to constrain in the current keyframe
            for jname in key_joint_names:
                idx = base_ee.index(jname)
                key_rot_names += base_rot_names[idx]
                key_pos_names += base_pos_names[idx]
            rot_joint_names.append(key_rot_names)
            pos_joint_names.append(key_pos_names)
        return rot_joint_names, pos_joint_names

    def __repr__(self):
        if self.folder is None:
            return f"{self.__class__.__name__}()"
        return f'{self.__class__.__name__}(folder="{self.folder}")'

    @property
    def device(self):
        """Device where neutral-joint buffers are stored.

        Returns 'cpu' if neutral_joints is not present.
        """
        if getattr(self, "neutral_joints", None) is None:
            return "cpu"
        return self.neutral_joints.device

    def fk(self, local_joint_rots: torch.Tensor, root_positions: torch.Tensor):
        """Run forward kinematics for this skeleton layout.

        Args:
            local_joint_rots: Local joint rotation matrices with shape
                `(..., J, 3, 3)`.
            root_positions: Root translations with shape `(..., 3)`.

        Returns:
            Tuple of `(global_joint_rots, posed_joints, posed_joints_norootpos)`.
        """
        global_joint_rots, posed_joints, posed_joints_norootpos = fk(local_joint_rots, root_positions, self)
        return global_joint_rots, posed_joints, posed_joints_norootpos

    def to_standard_tpose(self, local_rot_mats: torch.Tensor):
        """Convert local rotations into the skeleton's standard T-pose frame."""
        return to_standard_tpose(local_rot_mats, self)

    def from_standard_tpose(self, local_rot_mats: torch.Tensor):
        """Convert local rotations from the skeleton's standard T-pose frame."""
        return from_standard_tpose(local_rot_mats, self)

    def global_rots_to_local_rots(self, global_joint_rots: torch.Tensor):
        """Convert global joint rotations to local rotations for this hierarchy."""
        return global_rots_to_local_rots(global_joint_rots, self)

    def get_skel_slice(self, skeleton: "SkeletonBase"):
        """Build index mapping from another skeleton into this skeleton order.

        Args:
            skeleton: Source skeleton whose joint order is used by input tensors.

        Returns:
            A list of source indices ordered as `self.bone_order_names`.

        Raises:
            ValueError: If at least one required joint is missing from `skeleton`.
        """
        try:
            skel_slice = [skeleton.bone_index[x] for x in self.bone_order_names]
        except KeyError:
            raise ValueError("The current skeleton contain joints that are not in the input")
        return skel_slice