pymomentum.torch

pymomentum.torch.character

class pymomentum.torch.character.AdvancedInverseParameterTransform(parameter_transform: ParameterTransform, rcond: float = 1e-05)

Bases: InverseParameterTransform

This is a more advanced version of InverseParameterTransform that can handle the case where pose and scale are mutually exclusive (i.e. no joint parameter is affected by both pose and scale). In this case, we can compute a pose-only SVD and use that to compute the inverse parameter transform.

__init__(parameter_transform: ParameterTransform, rcond: float = 1e-05) → None: Initialize internal Module state, shared by both nn.Module and ScriptModule.

check_if_pose_scale_are_mutually_exclusive() → Tuple[bool, Tensor]

joint_parameters_to_model_parameters_with_known_scales(joint_parameters: Tensor, scale_parameters: Tensor) → Tensor

class pymomentum.torch.character.BlendShape(character: Character, *, dtype: dtype = torch.float32)

Bases: BlendShapeBase

__init__(character: Character, *, dtype: dtype = torch.float32) → None: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(coeffs: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class pymomentum.torch.character.BlendShapeBase(shape_vectors: Tensor, *, dtype: dtype = torch.float32)

Bases: Module

__init__(shape_vectors: Tensor, *, dtype: dtype = torch.float32) → None: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(coeffs: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class pymomentum.torch.character.Character(character: Character, *, has_parameter_transform: bool = True, has_skeleton: bool = True, has_rest_mesh: bool = True, has_skinning: bool = True, has_limits: bool = True, dtype: dtype = torch.float32)

Bases: Module

__init__(character: Character, *, has_parameter_transform: bool = True, has_skeleton: bool = True, has_rest_mesh: bool = True, has_skinning: bool = True, has_limits: bool = True, dtype: dtype = torch.float32) → None: Initialize internal Module state, shared by both nn.Module and ScriptModule.

bind_pose() → Tensor

joint_parameters_to_local_skeleton_state(joint_parameters: Tensor, backend: str = 'torch') → Tensor

joint_parameters_to_skeleton_state(joint_parameters: Tensor, use_double_precision: bool = True, backend: str = 'torch') → Tensor

local_skeleton_state_to_skeleton_state(local_skel_state: Tensor, use_double_precision: bool = True) → Tensor

model_parameters_to_blendshape_coefficients(model_parameters: Tensor) → Tensor

model_parameters_to_joint_parameters(model_parameters: Tensor) → Tensor

model_parameters_to_local_skeleton_state(model_parameters: Tensor, backend: str = 'torch') → Tensor

model_parameters_to_skeleton_state(model_parameters: Tensor, use_double_precision: bool = True, backend: str = 'torch') → Tensor

skin_points(skel_state: Tensor, rest_vertex_positions: Tensor | None = None, backend: str = 'torch') → Tensor

unpose(skel_state: Tensor, vertex_positions: Tensor) → Tensor

class pymomentum.torch.character.InverseParameterTransform(parameter_transform: ParameterTransform, rcond: float = 1e-05)

Bases: Module

__init__(parameter_transform: ParameterTransform, rcond: float = 1e-05) → None: Initialize internal Module state, shared by both nn.Module and ScriptModule.

joint_parameters_to_model_parameters(joint_parameters: Tensor) → Tensor

class pymomentum.torch.character.LinearBlendSkinning(character: Character, *, dtype: dtype = torch.float32)

Bases: Module

__init__(character: Character, *, dtype: dtype = torch.float32) → None: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(skel_state: Tensor, rest_vertex_positions: Tensor, backend: str = 'torch') → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

skin_with_transforms(global_state_t: Tensor, global_state_r: Tensor, global_state_s: Tensor, rest_vertex_positions: Tensor) → Tensor

Apply skinning using individual global TRS tensors (legacy interface).

This method provides backward compatibility for code that passes individual tensors instead of the TRS tuple format. Internally, it creates a tuple and delegates to skin_with_trs().

Args:: global_state_t: Global joint translations, shape (batch_size, num_joints, 3) global_state_r: Global joint rotations as 3x3 matrices, shape (batch_size, num_joints, 3, 3) global_state_s: Global joint scales, shape (batch_size, num_joints, 1) rest_vertex_positions: Rest vertex positions, shape (batch_size, num_vertices, 3) or (num_vertices, 3) which will be expanded to batch size
Returns:: skinned_vertices: Skinned vertex positions, shape (batch_size, num_vertices, 3)
Note:: This method is provided for backward compatibility. New code should use skin_with_trs() with the TRS tuple interface.

skin_with_trs(global_trs: Tuple[Tensor, Tensor, Tensor], rest_vertex_positions: Tensor) → Tensor

Apply skinning using global TRS state (Translation-Rotation-Scale).

This method takes global TRS state as a tuple and applies linear blend skinning using the TRS backend. This is useful when working directly with TRS representations or when you want to avoid conversions to skeleton state format.

Args:

global_trs: Global TRS transform tuple (translation, rotation_matrix, scale) where:

translation: Global joint translations, shape (batch_size, num_joints, 3)
rotation_matrix: Global joint rotations as 3x3 matrices, shape (batch_size, num_joints, 3, 3)
scale: Global joint scales, shape (batch_size, num_joints, 1)

rest_vertex_positions: Rest vertex positions, shape (batch_size, num_vertices, 3)

or (num_vertices, 3) which will be expanded to batch size

Returns:

skinned_vertices: Skinned vertex positions, shape (batch_size, num_vertices, 3)

Note:

This method uses the TRS backend skinning function for efficient computation.

unpose(global_joint_state: Tensor, verts: Tensor) → Tensor

class pymomentum.torch.character.Mesh(character: Character, *, dtype: dtype = torch.float32)

Bases: Module

__init__(character: Character, *, dtype: dtype = torch.float32) → None: Initialize internal Module state, shared by both nn.Module and ScriptModule.

class pymomentum.torch.character.ParameterTransform(character: Character, *, dtype: dtype = torch.float32)

Bases: Module

__init__(character: Character, *, dtype: dtype = torch.float32) → None: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(model_parameters: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class pymomentum.torch.character.Skeleton(character: Character, *, dtype: dtype = torch.float32)

Bases: Module

__init__(character: Character, *, dtype: dtype = torch.float32) → None: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(joint_parameters: Tensor, use_double_precision: bool = True) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

global_trs_to_local_trs(global_trs: Tuple[Tensor, Tensor, Tensor]) → Tuple[Tensor, Tensor, Tensor]

Convert global TRS state to local TRS state.

This function computes local joint transformations from global TRS transformations by multiplying each joint’s global transform by the inverse of its parent’s global transform. This is the TRS equivalent of skeleton_state_to_local_skeleton_state().

The computation is fully vectorized and parallel across all joints, making it efficient for GPU execution.

Args:

global_trs: Global TRS transform tuple (translation, rotation_matrix, scale) where:

translation: Global joint translations, shape (batch_size, num_joints, 3)
rotation_matrix: Global joint rotations as 3x3 matrices, shape (batch_size, num_joints, 3, 3)
scale: Global joint scales, shape (batch_size, num_joints, 1)

Returns:

TRS transform tuple (translation, rotation_matrix, scale) where:

translation: Local joint translations, shape (batch_size, num_joints, 3)
rotation_matrix: Local joint rotations as 3x3 matrices, shape (batch_size, num_joints, 3, 3)
scale: Local joint scales, shape (batch_size, num_joints, 1)

Note:

For root joints (parent index == -1), the global transform is returned as-is since they have no parent. This is the inverse operation of local_trs_to_global_trs().

joint_parameters_to_local_skeleton_state(joint_parameters: Tensor, backend: str = 'torch') → Tensor

joint_parameters_to_local_trs(joint_parameters: Tensor) → Tuple[Tensor, Tensor, Tensor]

Convert joint parameters to local TRS (Translation-Rotation-Scale) state.

This function takes joint parameters and converts them to the TRS representation as a tuple, consistent with the pymomentum.trs module interface. The TRS format is useful for applications that need explicit access to individual transformation components and can be more efficient than the skeleton state format.

Args:

joint_parameters: Joint parameters tensor of shape (batch_size, num_joints * 7): or (batch_size, num_joints, 7). Each joint has 7 parameters: - [0:3] translation offset - [3:6] euler xyz rotation angles - [6] log2 scale factor

Returns:

TRS transform tuple (translation, rotation_matrix, scale) where:

translation: Local joint translations, shape (batch_size, num_joints, 3)
rotation_matrix: Local joint rotations as 3x3 matrices, shape (batch_size, num_joints, 3, 3)
scale: Local joint scales, shape (batch_size, num_joints, 1)

Raises:

RuntimeError: If the character has no skeleton

Note:

This function is equivalent to joint_parameters_to_local_skeleton_state() but returns the TRS format instead of the 8-parameter skeleton state format. It uses the TRS backend for efficient computation.

joint_parameters_to_skeleton_state(joint_parameters: Tensor, use_double_precision: bool = True, backend: str = 'torch') → Tensor

joint_parameters_to_trs(joint_parameters: Tensor) → Tuple[Tensor, Tensor, Tensor]

Convert joint parameters directly to global TRS (Translation-Rotation-Scale) state.

This function performs the full forward kinematics pipeline using the TRS backend, converting joint parameters to local TRS state and then to global TRS state in a single call. This is more efficient than calling the two-step process separately.

Args:

joint_parameters: Joint parameters tensor of shape (batch_size, num_joints * 7): or (batch_size, num_joints, 7). Each joint has 7 parameters: - [0:3] translation offset - [3:6] euler xyz rotation angles - [6] log2 scale factor

Returns:

TRS transform tuple (translation, rotation_matrix, scale) where:

translation: Global joint translations, shape (batch_size, num_joints, 3)
rotation_matrix: Global joint rotations as 3x3 matrices, shape (batch_size, num_joints, 3, 3)
scale: Global joint scales, shape (batch_size, num_joints, 1)

Raises:

RuntimeError: If the character has no skeleton

Note:

This is equivalent to calling joint_parameters_to_local_trs() followed by local_trs_to_global_trs(), but more efficient as it avoids intermediate steps.

local_skeleton_state_to_joint_parameters(local_skel_state: Tensor) → Tensor

local_skeleton_state_to_skeleton_state(local_skel_state: Tensor, use_double_precision: bool = True) → Tensor

local_trs_to_global_trs(local_trs: Tuple[Tensor, Tensor, Tensor]) → Tuple[Tensor, Tensor, Tensor]

Convert local TRS state to global TRS state using forward kinematics.

This function performs forward kinematics on the local TRS transformations to compute the global joint transformations. This is useful when you have local TRS states from other sources and need to convert them to global states.

Args:

local_trs: Local TRS transform tuple (translation, rotation_matrix, scale) where:

translation: Local joint translations, shape (batch_size, num_joints, 3)
rotation_matrix: Local joint rotations as 3x3 matrices, shape (batch_size, num_joints, 3, 3)
scale: Local joint scales, shape (batch_size, num_joints, 1)

Returns:

TRS transform tuple (translation, rotation_matrix, scale) where:

translation: Global joint translations, shape (batch_size, num_joints, 3)
rotation_matrix: Global joint rotations as 3x3 matrices, shape (batch_size, num_joints, 3, 3)
scale: Global joint scales, shape (batch_size, num_joints, 1)

Note:

This function uses the TRS backend for efficient forward kinematics computation.

skeleton_state_to_joint_parameters(skel_state: Tensor, backend: str = 'torch') → Tensor

skeleton_state_to_local_skeleton_state(skel_state: Tensor) → Tensor

pymomentum.torch.parameter_limits

class pymomentum.torch.parameter_limits.ParameterLimits(parameter_limits: list[ParameterLimit], dtype: dtype = torch.float32)

Bases: Module

Native PyTorch implementation of momentum’s LimitErrorFunction, used to enforce parameter limits using a soft constraint.

__init__(parameter_limits: list[ParameterLimit], dtype: dtype = torch.float32) → None: Initialize internal Module state, shared by both nn.Module and ScriptModule.

evaluate_ellipsoid_error(skel_state: Tensor) → Tensor

evaluate_halfplane_error(model_params: Tensor) → Tensor

evaluate_linear_error(model_params: Tensor) → Tensor

evaluate_linear_joint_error(joint_params: Tensor) → Tensor

evaluate_minmax_error(model_params: Tensor) → Tensor

evaluate_minmaxjoint_error(joint_params: Tensor) → Tensor

forward(model_parameters: Tensor, joint_parameters: Tensor, skel_state: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

pymomentum.torch.sdf_collision

Pure PyTorch SDF collision geometry for torch_character.

Provides differentiable SDF evaluation matching the behavior of momentum.SDFColliderT. All operations use pure PyTorch ops, enabling gradient computation through autograd including higher-order gradients (gradient-of-gradient).

Typical usage:

import pymomentum.geometry as pym_geometry
from pymomentum.torch.sdf_collision import SDFCollider, SDFCollisionGeometry

# Create from pymomentum colliders
pym_colliders = [pym_geometry.SDFCollider(sdf=sdf, parent=joint_idx)]
collision_geo = SDFCollisionGeometry.from_pymomentum(pym_colliders)

# Evaluate: skel_state from Skeleton.forward(), points in world space
distances = collision_geo.evaluate(skel_state, world_points)

class pymomentum.torch.sdf_collision.SDFCollider(sdf_data: Tensor, bounds_min: Tensor, voxel_size: Tensor, local_translation: Tensor, local_rotation: Tensor, parent_joint: int = -1)

Bases: Module

A signed distance field volume attached to a skeleton joint.

Mirrors momentum.SDFColliderT with evaluation implemented in pure PyTorch for full autograd differentiability.

The SDF data is stored as a 3D tensor in (nz, ny, nx) (DHW) order: sdf_data[k, j, i] is the signed distance at grid point (i, j, k). This layout is compatible with F.grid_sample.

Parameters:

sdf_data – 3D SDF grid values, shape (nz, ny, nx) (DHW order).
bounds_min – SDF bounding box minimum corner, shape (3,).
voxel_size – Voxel dimensions in world units, shape (3,).
local_translation – Translation offset from parent joint, shape (3,).
local_rotation – Rotation matrix relative to parent joint, shape (3, 3).
parent_joint – Parent joint index. Use -1 for world-space SDFs.

__init__(sdf_data: Tensor, bounds_min: Tensor, voxel_size: Tensor, local_translation: Tensor, local_rotation: Tensor, parent_joint: int = -1) → None: Initialize internal Module state, shared by both nn.Module and ScriptModule.

evaluate(skel_state: Tensor, points: Tensor) → Tensor

Evaluate SDF distances for world-space query points.

Transforms points from world space into the SDF’s local coordinate system using the joint transform from skel_state, then evaluates via trilinear interpolation.

Parameters:

skel_state – Global skeleton state from forward kinematics, shape (..., n_joints, 8) with per-joint format [tx, ty, tz, qx, qy, qz, qw, scale]. Ignored when parent_joint == -1 (world-space SDF).
points – World-space query points, shape (..., n_points, 3).

Returns:

SDF distances at each query point, shape (..., n_points).

static from_pymomentum(collider: object, dtype: dtype = torch.float32) → SDFCollider

Create an SDFCollider from a pymomentum.geometry.SDFCollider.

Converts the SDF grid data to a (nz, ny, nx) PyTorch tensor (DHW order, ready for F.grid_sample) and the collider’s local transform to a rotation matrix.

Parameters:

collider – A pymomentum.geometry.SDFCollider instance.
dtype – Tensor dtype (default: torch.float32).

Returns:

A new SDFCollider module.

parent_joint: int

class pymomentum.torch.sdf_collision.SDFCollisionGeometry(colliders: list[SDFCollider])

Bases: Module

Collection of SDF colliders for character collision detection.

Mirrors momentum.SDFCollisionGeometryT — a list of SDFCollider modules attached to skeleton joints.

Parameters:: colliders – List of SDFCollider modules.

__init__(colliders: list[SDFCollider]) → None: Initialize internal Module state, shared by both nn.Module and ScriptModule.

evaluate(skel_state: Tensor, points: Tensor) → list[Tensor]

Evaluate all colliders and return per-collider SDF distances.

Parameters:

skel_state – Global skeleton state, shape (..., n_joints, 8).
points – World-space query points, shape (..., n_points, 3).

Returns:

List of distance tensors, one per collider, each shape (..., n_points).

evaluate_min(skel_state: Tensor, points: Tensor) → Tensor

Minimum SDF distance across all colliders for each point.

Useful for collision detection: negative values indicate penetration into at least one collider.

Parameters:

skel_state – Global skeleton state, shape (..., n_joints, 8).
points – World-space query points, shape (..., n_points, 3).

Returns:

Minimum signed distance per point, shape (..., n_points).

static from_pymomentum(colliders: list[object], dtype: dtype = torch.float32) → SDFCollisionGeometry

Create from a list of pymomentum.geometry.SDFCollider objects.

Parameters:

colliders – List of pymomentum.geometry.SDFCollider instances.
dtype – Tensor dtype (default: torch.float32).

Returns:

A new SDFCollisionGeometry module.

pymomentum.torch.sdf_collision.evaluate_sdf(sdf_data: Tensor, bounds_min: Tensor, voxel_size: Tensor, points: Tensor) → Tensor

Evaluate SDF distances at query points in SDF local space.

Performs world-to-grid conversion, boundary clamping, trilinear interpolation, and out-of-bounds distance computation, matching axel::SignedDistanceField::sample().

All operations are differentiable via PyTorch autograd. For out-of-bounds points, the Euclidean distance to the nearest grid boundary is added to the interpolated boundary value, providing correct gradient behavior:

Inside the grid: gradient comes from trilinear interpolation derivatives.
Outside the grid: gradient points away from the boundary (increasing distance).

Parameters:

sdf_data – 3D SDF grid values, shape (nz, ny, nx) (DHW order).
bounds_min – SDF bounding box minimum corner, shape (3,).
voxel_size – Size of each voxel in world units, shape (3,).
points – Query points in SDF local space, shape (..., 3).

Returns:

Signed distance values, shape (...).

pymomentum.torch

pymomentum.torch.character

pymomentum.torch.parameter_limits

pymomentum.torch.sdf_collision

pymomentum.torch.utility