fastdev.robo.warp_kinematics ============================ .. py:module:: fastdev.robo.warp_kinematics Module Contents --------------- .. py:function:: axis_angle_to_tf_mat(axis: warp.vec3, angle: warp.float32) .. py:function:: axis_distance_to_tf_mat(axis: warp.vec3, distance: warp.float32) .. py:function:: 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)) .. py:class:: ForwardKinematics Bases: :py:obj:`torch.autograd.Function` .. py:method:: 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] :staticmethod: .. py:method:: 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]]] :staticmethod: .. py:function:: 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] .. py:function:: 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] .. py:function:: 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. .. py:class:: CalculateJacobian Bases: :py:obj:`torch.autograd.Function` .. py:method:: 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]] :staticmethod: Forward pass to compute Jacobian and target link poses. .. py:method:: 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]]] :staticmethod: :abstractmethod: Backward pass for propagating Jacobian gradients. .. py:function:: 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]]] .. py:function:: 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. .. py:function:: 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] .. py:function:: 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. :rtype: y .. py:function:: inverse_lu(A: Any) -> Any Computes the inverse of a square matrix using LU decomposition without pivoting. :returns: the inverse of A. :rtype: A_inv .. py:function:: 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. .. py:function:: 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] .. py:function:: 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]]