pymomentum-cpu 0.1.77.post26__cp312-cp312-macosx_15_0_arm64.whl

pymomentum/skel_state.py

@@ -0,0 +1,514 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Skeleton State Utilities
+========================
+
+This module provides utilities for working with skeleton states in PyMomentum.
+
+A skeleton state is a compact 8-dimensional representation of a 3D transformation
+that encodes translation, rotation, and uniform scale. The 8 components are organized as:
+
+- **Translation (3 components)**: tx, ty, tz - 3D position offset
+- **Rotation (4 components)**: rx, ry, rz, rw - quaternion rotation (x, y, z, w)
+- **Scale (1 component)**: s - uniform scale factor
+
+This representation is generated by forward kinematics functions like
+:func:`pymomentum.geometry.joint_parameters_to_skeleton_state` and represents the
+transformation from a local joint space to the world space.
+
+- Creating skeleton states from individual components (:func:`from_translation`,
+  :func:`from_quaternion`, :func:`from_scale`)
+- Converting between skeleton states and 4x4 transformation matrices
+  (:func:`to_matrix`, :func:`from_matrix`)
+- Performing transformations and operations (:func:`multiply`, :func:`inverse`,
+  :func:`transform_points`)
+- Interpolation and blending (:func:`slerp`, :func:`blend`)
+- Validation and manipulation (:func:`check`, :func:`split`)
+
+Example:
+    Creating and using a skeleton state::
+
+        import torch
+        from pymomentum import skel_state
+
+        # Create identity transform
+        identity = skel_state.identity()
+
+        # Create from translation
+        translation = torch.tensor([1.0, 2.0, 3.0])
+        state = skel_state.from_translation(translation)
+
+        # Transform points
+        points = torch.tensor([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])
+        transformed = skel_state.transform_points(state, points)
+
+Note:
+    All functions in this module expect skeleton states to be PyTorch tensors
+    with the last dimension having size 8, following the (tx, ty, tz, rx, ry, rz, rw, s)
+    format.
+"""
+
+from typing import Sequence
+
+import torch
+from pymomentum import quaternion
+
+# pyre-strict
+
+
+def check(skel_state: torch.Tensor) -> None:
+    """
+    Check if the skeleton state has the correct shape.
+
+    :parameter skel_state: The skeleton state to check.
+    :type skel_state: torch.Tensor
+    :raises ValueError: If the skeleton state does not have the correct shape.
+    """
+    if skel_state.shape[-1] != 8:
+        raise ValueError(
+            "Expected skeleton state to have last dimension 8 (tx, ty, tz, rx, ry, rz, rw, s)"
+        )
+
+
+def split(skel_state: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Split a skeleton state into translation, rotation, and scale components.
+
+    :parameter skel_state: The skeleton state to split.
+    :type skel_state: torch.Tensor
+    :return: A tuple of tensors (translation, rotation, scale).
+    :rtype: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+    """
+    check(skel_state)
+    return skel_state[..., :3], skel_state[..., 3:7], skel_state[..., 7:]
+
+
+def from_translation(translation: torch.Tensor) -> torch.Tensor:
+    """
+    Create a skeleton state from translation.
+
+    :parameter translation: The translation component.
+    :type translation: torch.Tensor
+    :return: The skeleton state.
+    :rtype: torch.Tensor
+    """
+    return torch.cat(
+        (
+            translation,
+            quaternion.identity(
+                size=translation.shape[:-1],
+                device=translation.device,
+                dtype=translation.dtype,
+            ),
+            torch.ones(
+                *translation.shape[:-1],
+                1,
+                dtype=translation.dtype,
+                device=translation.device,
+            ),
+        ),
+        -1,
+    )
+
+
+def from_quaternion(rotation: torch.Tensor) -> torch.Tensor:
+    """
+    Create a skeleton state from rotation.
+
+    :parameter rotation: The rotation component.
+    :type rotation: torch.Tensor
+    :return: The skeleton state.
+    :rtype: torch.Tensor
+    """
+    return torch.cat(
+        (
+            torch.zeros(
+                *rotation.shape[:-1],
+                3,
+                dtype=rotation.dtype,
+                device=rotation.device,
+            ),
+            rotation,
+            torch.ones(
+                *rotation.shape[:-1], 1, dtype=rotation.dtype, device=rotation.device
+            ),
+        ),
+        -1,
+    )
+
+
+def from_scale(scale: torch.Tensor) -> torch.Tensor:
+    """
+    Create a skeleton state from scale.
+
+    :parameter scale: The scale component.
+    :type scale: torch.Tensor
+    :return: The skeleton state.
+    :rtype: torch.Tensor
+    """
+    return torch.cat(
+        (
+            torch.zeros(
+                *scale.shape[:-1],
+                3,
+                dtype=scale.dtype,
+                device=scale.device,
+            ),
+            quaternion.identity(
+                size=scale.shape[:-1],
+                device=scale.device,
+                dtype=scale.dtype,
+            ),
+            scale,
+        ),
+        -1,
+    )
+
+
+def to_matrix(skeleton_state: torch.Tensor) -> torch.Tensor:
+    """
+    Convert skeleton state to a tensor of 4x4 matrices. The matrix represents the transform from a local joint space to the world space.
+
+    :parameter skeleton_state: The skeleton state to convert.
+    :type skeleton_state: torch.Tensor
+    :return: A tensor containing 4x4 matrix transforms.
+    :rtype: torch.Tensor
+    """
+    check(skeleton_state)
+    t, q, s = split(skeleton_state)
+
+    # Assuming quaternionToRotationMatrix is implemented elsewhere
+    rot_mat = quaternion.to_rotation_matrix(q)
+    linear = rot_mat * s.unsqueeze(-2).expand_as(rot_mat)
+    affine = torch.cat((linear, t.unsqueeze(-1)), -1)
+
+    last_row = (
+        torch.tensor(
+            [0.0, 0.0, 0.0, 1.0],
+            device=skeleton_state.device,
+            dtype=skeleton_state.dtype,
+        )
+        .unsqueeze(0)
+        .expand(*affine.shape[:-2], 1, 4)
+    )
+    result = torch.cat((affine, last_row), -2)
+    return result
+
+
+def _normalize_split_skel_state(
+    components: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Normalize the quaternion component of a split skeleton state.
+
+    :parameter components: Tuple of (translation, rotation, scale) components.
+    :type components: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+    :return: Tuple of (translation, normalized_rotation, scale) components.
+    :rtype: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+    """
+    t, q, s = components
+    return t, quaternion.normalize(q), s
+
+
+def _multiply_split_skel_states(
+    skel_state1: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+    skel_state2: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+) -> torch.Tensor:
+    """
+    Helper function to multiply two skeleton states from their split components.
+    Assumes quaternions are already normalized.
+
+    :parameter skel_state1: Tuple of (translation, rotation, scale) for first skeleton state.
+    :type skel_state1: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+    :parameter skel_state2: Tuple of (translation, rotation, scale) for second skeleton state.
+    :type skel_state2: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+    :return: The product of the two skeleton states.
+    :rtype: torch.Tensor
+    """
+    t1, q1, s1 = skel_state1
+    t2, q2, s2 = skel_state2
+
+    t_res = t1 + s1 * quaternion.rotate_vector_assume_normalized(q1, t2)
+    s_res = s1 * s2
+    q_res = quaternion.multiply_assume_normalized(q1, q2)
+
+    return torch.cat((t_res, q_res, s_res), -1)
+
+
+def multiply(s1: torch.Tensor, s2: torch.Tensor) -> torch.Tensor:
+    """
+    Multiply two skeleton states.
+
+    :parameter s1: The first skeleton state.
+    :type s1: torch.Tensor
+    :parameter s2: The second skeleton state.
+    :type s2: torch.Tensor
+    :return: The product of the two skeleton states.
+    :rtype: torch.Tensor
+    """
+    check(s1)
+    check(s2)
+    while s1.ndim < s2.ndim:
+        s1 = s1.unsqueeze(0)
+
+    return _multiply_split_skel_states(
+        _normalize_split_skel_state(split(s1)), _normalize_split_skel_state(split(s2))
+    )
+
+
+def multiply_assume_normalized(s1: torch.Tensor, s2: torch.Tensor) -> torch.Tensor:
+    """
+    Multiply two skeleton states.
+    This is an optimized version that assumes both skeleton states are already properly
+    formatted and the quaternions are normalized, skipping validation checks.
+
+    :parameter s1: The first skeleton state.
+    :type s1: torch.Tensor
+    :parameter s2: The second skeleton state.
+    :type s2: torch.Tensor
+    :return: The product of the two skeleton states.
+    :rtype: torch.Tensor
+    """
+    while s1.ndim < s2.ndim:
+        s1 = s1.unsqueeze(0)
+
+    return _multiply_split_skel_states(split(s1), split(s2))
+
+
+def inverse(skeleton_states: torch.Tensor) -> torch.Tensor:
+    """
+    Compute the inverse of a skeleton state.
+
+    :parameter skeleton_states: The skeleton state to invert.
+    :type skeleton_states: torch.Tensor
+    :return: The inverted skeleton state.
+    :rtype: torch.Tensor
+    """
+    t, q, s = split(skeleton_states)
+    q_inv = quaternion.inverse(q)
+    s_inv = torch.reciprocal(s)
+
+    return torch.cat((-s_inv * quaternion.rotate_vector(q_inv, t), q_inv, s_inv), -1)
+
+
+def _transform_points_split_skel_state(
+    skel_state: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+    points: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Helper function to transform points using skeleton state components.
+    Assumes quaternion is already normalized.
+
+    :parameter skel_state: Tuple of (translation, rotation, scale) components.
+    :type skel_state: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+    :parameter points: The points to transform.
+    :type points: torch.Tensor
+    :return: The transformed points.
+    :rtype: torch.Tensor
+    """
+    t, q, s = skel_state
+    return t + quaternion.rotate_vector(q, s * points)
+
+
+def transform_points(skel_state: torch.Tensor, points: torch.Tensor) -> torch.Tensor:
+    """
+    Transform 3d points by the transform represented by the skeleton state.
+
+    :parameter skel_state: The skeleton state to use for transformation.
+    :type skel_state: torch.Tensor
+    :parameter points: The points to transform.
+    :type points: torch.Tensor
+    :return: The transformed points.
+    :rtype: torch.Tensor
+    """
+    check(skel_state)
+    if points.dim() < 1 or points.shape[-1] != 3:
+        raise ValueError("Points tensor should have last dimension 3.")
+
+    # allow transforming multiple points with a single skel_state.
+    while skel_state.ndim < points.ndim:
+        skel_state = skel_state.unsqueeze(0)
+
+    return _transform_points_split_skel_state(
+        _normalize_split_skel_state(split(skel_state)), points
+    )
+
+
+def transform_points_assume_normalized(
+    skel_state: torch.Tensor, points: torch.Tensor
+) -> torch.Tensor:
+    """
+    Transform 3d points by the transform represented by the skeleton state.
+    This is an optimized version that assumes the skeleton state is already properly
+    formatted and the quaternion is normalized, skipping validation checks.
+
+    :parameter skel_state: The skeleton state to use for transformation.
+    :type skel_state: torch.Tensor
+    :parameter points: The points to transform.
+    :type points: torch.Tensor
+    :return: The transformed points.
+    :rtype: torch.Tensor
+    """
+    # when multipying by points, we'll allow a single skel_state to transform multiple points.  So we'll allow adding a single extra dimension to the skel_state.
+    if skel_state.ndim < points.ndim:
+        skel_state = skel_state.unsqueeze(0)
+    return _transform_points_split_skel_state(split(skel_state), points)
+
+
+def identity(
+    size: Sequence[int] | None = None, device: torch.device | None = None
+) -> torch.Tensor:
+    """
+    Returns a skeleton state representing the identity transform.
+
+    :parameter sizes: The size of each dimension in the output tensor. Defaults to None, which means the output will be a 1D tensor with 8 elements.
+    :type sizes: list[int], optional
+    :parameter device: The device on which to create the tensor. Defaults to None, which means the tensor will be created on the default device.
+    :type device: torch.device, optional
+    :return: The identity skeleton state.
+    :rtype: torch.Tensor
+    """
+    zeros = (
+        torch.zeros(*size, 3, device=device) if size else torch.zeros(3, device=device)
+    )
+    ones = torch.ones(*size, 1, device=device) if size else torch.ones(1, device=device)
+    q_identity = quaternion.identity(size=size, device=device)
+
+    return torch.cat((zeros, q_identity, ones), -1)
+
+
+def blend(
+    skel_states: torch.Tensor, weights: torch.Tensor | None = None
+) -> torch.Tensor:
+    """
+    Blend k skeleton states with the passed-in weights.
+
+    :parameter skel_states: The skeleton states to blend.
+    :type skel_states: torch.Tensor
+    :parameter weights: The weights to use, if not provided, weights are assumed to be all 1s
+    :return: The blended skeleton state.
+    """
+    t, q, s = split(skel_states)
+    weights = quaternion.check_and_normalize_weights(q, weights)
+    t_blend = (weights.unsqueeze(-1).expand_as(t) * t).sum(-2)
+    q_blend = quaternion.blend(q, weights)
+    s_blend = (weights.unsqueeze(-1).expand_as(s) * s).sum(-2)
+    return torch.cat((t_blend, q_blend, s_blend), -1)
+
+
+def slerp(s0: torch.Tensor, s1: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    """
+    Spherical linear interpolation between two skeleton states.
+
+    :parameter s0: The first skeleton state.
+    :parameter s1: The second skeleton state.
+    :parameter t: The interpolation factor, where 0 <= t <= 1.  t=0 corresponds to s0, t=1 corresponds to s1.
+    :return: The interpolated skeleton state.
+    """
+    check(s0)
+    check(s1)
+
+    if t.ndim < s0.ndim:
+        t = t.unsqueeze(0)
+
+    t0, q0, s0 = split(s0)
+    t1, q1, s1 = split(s1)
+
+    s = (1 - t) * s0 + t * s1
+    q = quaternion.slerp(q0, q1, t)
+    t = (1 - t) * t0 + t * t1
+    return torch.cat((t, q, s), -1)
+
+
+def from_matrix(matrices: torch.Tensor) -> torch.Tensor:
+    """
+    Convert 4x4 matrices to skeleton states.  Assumes that the scale is uniform.
+
+    Args:
+        matrices (torch.Tensor): A tensor of 4x4 matrices.
+
+    Returns:
+        torch.Tensor: The corresponding skeleton states.
+    """
+    if matrices.dim() < 2 or matrices.shape[-1] != 4 or matrices.shape[-2] != 4:
+        raise ValueError("Expected a tensor of 4x4 matrices")
+    initial_shape = matrices.shape
+    if matrices.dim() == 2:
+        matrices = matrices.unsqueeze(0)
+    else:
+        matrices = matrices.flatten(0, -3)
+    linear = matrices[..., :3, :3]
+    translations = matrices[..., :3, 3]
+    U, S, Vt = torch.linalg.svd(linear)
+    scales = S[..., :1]
+    rotation_matrices = torch.bmm(U, Vt)
+    quaternions = quaternion.from_rotation_matrix(rotation_matrices)
+    result = torch.cat((translations, quaternions, scales), -1)
+    result_shape = list(initial_shape[:-2]) + [8]
+    return result.reshape(result_shape)
+
+
+@torch.jit.script
+def multiply_backprop(
+    state1: torch.Tensor, state2: torch.Tensor, grad_state: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Backpropagate gradients through skeleton state multiplication.
+
+    This function computes the gradients with respect to the two input skeleton states
+    given the gradient with respect to the output of skeleton state multiplication.
+
+    The forward formula is (matching the corrected skel_state.multiply):
+    - t = t1 + s1 * quaternion.rotate_vector(q1, t2)
+    - q = quaternion.multiply(q1, q2)
+    - s = s1 * s2
+
+    :parameter state1: The first skeleton state (parent).
+    :type state1: torch.Tensor
+    :parameter state2: The second skeleton state (child).
+    :type state2: torch.Tensor
+    :parameter grad_state: The gradient with respect to the output skeleton state.
+    :type grad_state: torch.Tensor
+    :return: A tuple of (grad_state1, grad_state2) representing gradients with respect
+             to the first and second skeleton states respectively.
+    :rtype: tuple[torch.Tensor, torch.Tensor]
+    """
+    with torch.no_grad():
+        t1, q1, s1 = split(state1)
+        t2, q2, s2 = split(state2)
+        gt, gq, gs = split(grad_state)
+
+        # Gradient with respect to s2: gs * s1
+        gs2 = gs * s1
+
+        # Gradient with respect to quaternions from quaternion multiplication
+        # Use the "assume normalized" version to avoid norm computation
+        gq1_from_gq, gq2 = quaternion.multiply_backprop_assume_normalized(q1, q2, gq)
+
+        # Gradient with respect to quaternions from translation computation
+        # Following the corrected forward: t = t1 + s1 * quaternion.rotate_vector(q1, t2)
+        # We need gradients from: quaternion.rotate_vector(q1, t2) with gt scaled by s1
+        gq1_from_gt, gt2 = quaternion.rotate_vector_backprop_assume_normalized(
+            q1, t2, gt
+        )
+        gt2 = gt2 * s1.expand_as(t2)
+
+        # Combine quaternion gradients - add the missing s1 factor!
+        gq1 = gq1_from_gq + gq1_from_gt * s1.expand_as(gq1_from_gt)
+
+        # Gradient with respect to t1: direct passthrough
+        gt1 = gt
+
+        # Gradient with respect to s1: from both scale multiplication and rotation
+        # From scale multiplication: gs * s2
+        # From rotation: sum over the dot product of rotated vector with gradient
+        gs1_from_rotation = (quaternion.rotate_vector(q1, t2) * gt).sum(
+            dim=-1, keepdim=True
+        )
+        gs1 = gs1_from_rotation + gs * s2
+
+    return torch.cat([gt1, gq1, gs1], dim=-1), torch.cat([gt2, gq2, gs2], dim=-1)