fastdev.robo.warp_kinematics

Module Contents

fastdev.robo.warp_kinematics.axis_angle_to_tf_mat(axis: warp.vec3, angle: warp.float32)

Parameters:

• axis (warp.vec3)

• angle (warp.float32)

fastdev.robo.warp_kinematics.axis_distance_to_tf_mat(axis: warp.vec3, distance: warp.float32)

Parameters:

• axis (warp.vec3)

• distance (warp.float32)

fastdev.robo.warp_kinematics.forward_kinematics_kernel(joint_values: wp.array2d(dtype=wp.float32), root_poses: wp.array2d(dtype=wp.mat44), joint_first_indices: wp.array(dtype=wp.int32), link_indices_topological_order: wp.array(dtype=wp.int32), parent_link_indices: wp.array(dtype=wp.int32), link_joint_indices: wp.array(dtype=wp.int32), link_joint_types: wp.array(dtype=wp.int32), link_joint_origins: wp.array(dtype=wp.mat44), link_joint_axes: wp.array(dtype=wp.vec3), link_first_indices: wp.array(dtype=wp.int32), link_poses: wp.array2d(dtype=wp.mat44))

Parameters:

• joint_values (wp.array2d(dtype=wp.float32))

• root_poses (wp.array2d(dtype=wp.mat44))

• joint_first_indices (wp.array(dtype=wp.int32))

• link_indices_topological_order (wp.array(dtype=wp.int32))

• parent_link_indices (wp.array(dtype=wp.int32))

• link_joint_indices (wp.array(dtype=wp.int32))

• link_joint_types (wp.array(dtype=wp.int32))

• link_joint_origins (wp.array(dtype=wp.mat44))

• link_joint_axes (wp.array(dtype=wp.vec3))

• link_first_indices (wp.array(dtype=wp.int32))

• link_poses (wp.array2d(dtype=wp.mat44))

class fastdev.robo.warp_kinematics.ForwardKinematics

Bases: torch.autograd.Function

static forward(ctx, joint_values: jaxtyping.Float[torch.Tensor, ... total_num_joints], articulation: fastdev.robo.articulation.Articulation, root_poses: Optional[jaxtyping.Float[torch.Tensor, ... num_arti 4 4]] = None) → jaxtyping.Float[torch.Tensor, ... total_num_links 4 4]

Parameters:

• joint_values (jaxtyping.Float[torch.Tensor, ... total_num_joints])

• articulation (fastdev.robo.articulation.Articulation)

• root_poses (Optional[jaxtyping.Float[torch.Tensor, ... num_arti 4 4]])

Return type:

jaxtyping.Float[torch.Tensor, … total_num_links 4 4]

static backward(ctx, link_poses_grad: jaxtyping.Float[torch.Tensor, ... total_num_links 4 4]) → Tuple[Optional[jaxtyping.Float[torch.Tensor, ... total_num_joints]], None, Optional[jaxtyping.Float[torch.Tensor, ... num_arti 4 4]]]

Parameters:

link_poses_grad (jaxtyping.Float[torch.Tensor, ... total_num_links 4 4])

Return type:

Tuple[Optional[jaxtyping.Float[torch.Tensor, … total_num_joints]], None, Optional[jaxtyping.Float[torch.Tensor, … num_arti 4 4]]]

fastdev.robo.warp_kinematics.forward_kinematics(joint_values: jaxtyping.Float[torch.Tensor, ... total_num_joints], articulation: fastdev.robo.articulation.Articulation, root_poses: Optional[jaxtyping.Float[torch.Tensor, ... num_arti 4 4]] = None) → jaxtyping.Float[torch.Tensor, ... total_num_links 4 4]

Parameters:

• joint_values (jaxtyping.Float[torch.Tensor, ... total_num_joints])

• articulation (fastdev.robo.articulation.Articulation)

• root_poses (Optional[jaxtyping.Float[torch.Tensor, ... num_arti 4 4]])

Return type:

jaxtyping.Float[torch.Tensor, … total_num_links 4 4]

fastdev.robo.warp_kinematics.forward_kinematics_numpy(joint_values: jaxtyping.Float[numpy.ndarray, ... total_num_joints], articulation: fastdev.robo.articulation.Articulation, root_poses: Optional[jaxtyping.Float[numpy.ndarray, ... num_arti 4 4]] = None) → jaxtyping.Float[numpy.ndarray, ... total_num_links 4 4]

Parameters:

• joint_values (jaxtyping.Float[numpy.ndarray, ... total_num_joints])

• articulation (fastdev.robo.articulation.Articulation)

• root_poses (Optional[jaxtyping.Float[numpy.ndarray, ... num_arti 4 4]])

Return type:

jaxtyping.Float[numpy.ndarray, … total_num_links 4 4]

fastdev.robo.warp_kinematics.calculate_jacobian_kernel(link_poses: wp.array2d(dtype=wp.mat44), target_link_indices: wp.array(dtype=wp.int32), ancestor_mask: wp.array(dtype=wp.int32), joint_axes: wp.array(dtype=wp.vec3), link_joint_types: wp.array(dtype=wp.int32), link_joint_indices: wp.array(dtype=wp.int32), link_joint_origins: wp.array(dtype=wp.mat44), parent_link_indices: wp.array(dtype=wp.int32), joint_first_indices: wp.array(dtype=wp.int32), link_first_indices: wp.array(dtype=wp.int32), jacobian: wp.array3d(dtype=wp.float32))

Compute the Jacobian using precomputed target link poses.

Parameters:

• link_poses (wp.array2d(dtype=wp.mat44))

• target_link_indices (wp.array(dtype=wp.int32))

• ancestor_mask (wp.array(dtype=wp.int32))

• joint_axes (wp.array(dtype=wp.vec3))

• link_joint_types (wp.array(dtype=wp.int32))

• link_joint_indices (wp.array(dtype=wp.int32))

• link_joint_origins (wp.array(dtype=wp.mat44))

• parent_link_indices (wp.array(dtype=wp.int32))

• joint_first_indices (wp.array(dtype=wp.int32))

• link_first_indices (wp.array(dtype=wp.int32))

• jacobian (wp.array3d(dtype=wp.float32))

class fastdev.robo.warp_kinematics.CalculateJacobian

Bases: torch.autograd.Function

static forward(ctx, joint_values: jaxtyping.Float[torch.Tensor, ... total_num_joints], target_link_indices: torch.Tensor, articulation: fastdev.robo.articulation.Articulation, root_poses: Optional[jaxtyping.Float[torch.Tensor, ... num_arti 4 4]] = None) → Tuple[jaxtyping.Float[torch.Tensor, ... 6 total_num_joints], jaxtyping.Float[torch.Tensor, ... num_arti 4 4]]

Forward pass to compute Jacobian and target link poses.

Parameters:

• joint_values (jaxtyping.Float[torch.Tensor, ... total_num_joints])

• target_link_indices (torch.Tensor)

• articulation (fastdev.robo.articulation.Articulation)

• root_poses (Optional[jaxtyping.Float[torch.Tensor, ... num_arti 4 4]])

Return type:

Tuple[jaxtyping.Float[torch.Tensor, … 6 total_num_joints], jaxtyping.Float[torch.Tensor, … num_arti 4 4]]

static backward(ctx, jacobian_grad: jaxtyping.Float[torch.Tensor, ... 6 total_num_joints], target_link_poses_grad: jaxtyping.Float[torch.Tensor, ... num_arti 4 4]) → Tuple[Optional[jaxtyping.Float[torch.Tensor, ... total_num_joints]], None, None, Optional[jaxtyping.Float[torch.Tensor, ... num_arti 4 4]]]

Abstractmethod:

Parameters:

• jacobian_grad (jaxtyping.Float[torch.Tensor, ... 6 total_num_joints])

• target_link_poses_grad (jaxtyping.Float[torch.Tensor, ... num_arti 4 4])

Return type:

Tuple[Optional[jaxtyping.Float[torch.Tensor, … total_num_joints]], None, None, Optional[jaxtyping.Float[torch.Tensor, … num_arti 4 4]]]

Backward pass for propagating Jacobian gradients.

fastdev.robo.warp_kinematics.calculate_jacobian(joint_values: jaxtyping.Float[torch.Tensor, ... total_num_joints], target_link_indices: torch.Tensor, articulation: fastdev.robo.articulation.Articulation, root_poses: Optional[jaxtyping.Float[torch.Tensor, ... num_arti 4 4]] = None, return_target_link_poses: bool = False) → Union[jaxtyping.Float[torch.Tensor, ... 6 total_num_joints], Tuple[jaxtyping.Float[torch.Tensor, ... 6 total_num_joints], jaxtyping.Float[torch.Tensor, ... num_arti 4 4]]]

Parameters:

• joint_values (jaxtyping.Float[torch.Tensor, ... total_num_joints])

• target_link_indices (torch.Tensor)

• articulation (fastdev.robo.articulation.Articulation)

• root_poses (Optional[jaxtyping.Float[torch.Tensor, ... num_arti 4 4]])

• return_target_link_poses (bool)

Return type:

Union[jaxtyping.Float[torch.Tensor, … 6 total_num_joints], Tuple[jaxtyping.Float[torch.Tensor, … 6 total_num_joints], jaxtyping.Float[torch.Tensor, … num_arti 4 4]]]

fastdev.robo.warp_kinematics.delta_pose_kernel(T_current: wp.array(dtype=wp.mat44), pos_target: wp.array(dtype=wp.vec3), quat_target: wp.array(dtype=wp.vec4), delta: wp.array2d(dtype=wp.float32)) → None

Compute delta pose error between current and target poses in axis-angle form.

Parameters:

• T_current (wp.array(dtype=wp.mat44))

• pos_target (wp.array(dtype=wp.vec3))

• quat_target (wp.array(dtype=wp.vec4))

• delta (wp.array2d(dtype=wp.float32))

Return type:

None

fastdev.robo.warp_kinematics.delta_pose_warp(T_current: jaxtyping.Float[torch.Tensor, ... 4 4], pos_target: jaxtyping.Float[torch.Tensor, ... 3], quat_target: jaxtyping.Float[torch.Tensor, ... 4]) → jaxtyping.Float[torch.Tensor, ... 6]

Parameters:

• T_current (jaxtyping.Float[torch.Tensor, ... 4 4])

• pos_target (jaxtyping.Float[torch.Tensor, ... 3])

• quat_target (jaxtyping.Float[torch.Tensor, ... 4])

Return type:

jaxtyping.Float[torch.Tensor, … 6]

fastdev.robo.warp_kinematics.solve_lower_triangular(L: Any, b: Any) → Any

Solves for y in L y = b where L is lower triangular with unit diagonal.

Returns:

the solution vector.

Return type:

y

Parameters:

• L (Any)

• b (Any)

fastdev.robo.warp_kinematics.inverse_lu(A: Any) → Any

Computes the inverse of a square matrix using LU decomposition without pivoting.

Returns:

the inverse of A.

Return type:

A_inv

Parameters:

A (Any)

fastdev.robo.warp_kinematics.compute_dq_kernel_v2(J: wp.array3d(dtype=wp.float32), err: wp.array3d(dtype=wp.float32), A: wp.array3d(dtype=wp.float32), joint_first_indices: wp.array(dtype=wp.int32), damping: warp.float32, dq: wp.array2d(dtype=wp.float32)) → None

Compute joint velocity dq using a damped least-squares formulation.

Parameters:

• J (wp.array3d(dtype=wp.float32))

• err (wp.array3d(dtype=wp.float32))

• A (wp.array3d(dtype=wp.float32))

• joint_first_indices (wp.array(dtype=wp.int32))

• damping (warp.float32)

• dq (wp.array2d(dtype=wp.float32))

Return type:

None

fastdev.robo.warp_kinematics.compute_dq_warp(J: jaxtyping.Float[torch.Tensor, ... 6 total_num_joints], err: jaxtyping.Float[torch.Tensor, ... num_arti 6], joint_first_indices_pt: jaxtyping.Int[torch.Tensor, compute_dq_warp.num_arti], damping: float) → jaxtyping.Float[torch.Tensor, ... total_num_joints]

Parameters:

• J (jaxtyping.Float[torch.Tensor, ... 6 total_num_joints])

• err (jaxtyping.Float[torch.Tensor, ... num_arti 6])

• joint_first_indices_pt (jaxtyping.Int[torch.Tensor, compute_dq_warp.num_arti])

• damping (float)

Return type:

jaxtyping.Float[torch.Tensor, … total_num_joints]

fastdev.robo.warp_kinematics.inverse_kinematics(target_link_poses: jaxtyping.Float[torch.Tensor, ... num_arti 4 4], target_link_indices: jaxtyping.Int[torch.Tensor, inverse_kinematics.num_arti], articulation: fastdev.robo.articulation.Articulation, max_iterations: int = 100, learning_rate: float = 0.1, tolerance: float = 1e-06, damping: float = 0.01, num_retries: int = 10, init_joint_values: Optional[jaxtyping.Float[torch.Tensor, ... total_num_dofs]] = None, jitter_strength: float = 1.0) → Tuple[jaxtyping.Float[torch.Tensor, ... total_num_dofs], jaxtyping.Bool[torch.Tensor, ... num_arti]]

Parameters:

• target_link_poses (jaxtyping.Float[torch.Tensor, ... num_arti 4 4])

• target_link_indices (jaxtyping.Int[torch.Tensor, inverse_kinematics.num_arti])

• articulation (fastdev.robo.articulation.Articulation)

• max_iterations (int)

• learning_rate (float)

• tolerance (float)

• damping (float)

• num_retries (int)

• init_joint_values (Optional[jaxtyping.Float[torch.Tensor, ... total_num_dofs]])

• jitter_strength (float)

Return type:

Tuple[jaxtyping.Float[torch.Tensor, … total_num_dofs], jaxtyping.Bool[torch.Tensor, … num_arti]]