warp-lang 1.9.1__py3-none-manylinux_2_34_aarch64.whl → 1.10.0rc2__py3-none-manylinux_2_34_aarch64.whl

warp/sim/articulation.py

@@ -1,793 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import warp as wp
-
-from .utils import quat_decompose, quat_twist
-
-
-@wp.func
-def compute_2d_rotational_dofs(
-    axis_0: wp.vec3,
-    axis_1: wp.vec3,
-    q0: float,
-    q1: float,
-    qd0: float,
-    qd1: float,
-):
-    """
-    Computes the rotation quaternion and 3D angular velocity given the joint axes, coordinates and velocities.
-    """
-    q_off = wp.quat_from_matrix(wp.matrix_from_cols(axis_0, axis_1, wp.cross(axis_0, axis_1)))
-
-    # body local axes
-    local_0 = wp.quat_rotate(q_off, wp.vec3(1.0, 0.0, 0.0))
-    local_1 = wp.quat_rotate(q_off, wp.vec3(0.0, 1.0, 0.0))
-
-    axis_0 = local_0
-    q_0 = wp.quat_from_axis_angle(axis_0, q0)
-
-    axis_1 = wp.quat_rotate(q_0, local_1)
-    q_1 = wp.quat_from_axis_angle(axis_1, q1)
-
-    rot = q_1 * q_0
-
-    vel = axis_0 * qd0 + axis_1 * qd1
-
-    return rot, vel
-
-
-@wp.func
-def invert_2d_rotational_dofs(
-    axis_0: wp.vec3,
-    axis_1: wp.vec3,
-    q_p: wp.quat,
-    q_c: wp.quat,
-    w_err: wp.vec3,
-):
-    """
-    Computes generalized joint position and velocity coordinates for a 2D rotational joint given the joint axes, relative orientations and angular velocity differences between the two bodies the joint connects.
-    """
-    q_off = wp.quat_from_matrix(wp.matrix_from_cols(axis_0, axis_1, wp.cross(axis_0, axis_1)))
-    q_pc = wp.quat_inverse(q_off) * wp.quat_inverse(q_p) * q_c * q_off
-
-    # decompose to a compound rotation each axis
-    angles = quat_decompose(q_pc)
-
-    # find rotation axes
-    local_0 = wp.quat_rotate(q_off, wp.vec3(1.0, 0.0, 0.0))
-    local_1 = wp.quat_rotate(q_off, wp.vec3(0.0, 1.0, 0.0))
-    local_2 = wp.quat_rotate(q_off, wp.vec3(0.0, 0.0, 1.0))
-
-    axis_0 = local_0
-    q_0 = wp.quat_from_axis_angle(axis_0, angles[0])
-
-    axis_1 = wp.quat_rotate(q_0, local_1)
-    q_1 = wp.quat_from_axis_angle(axis_1, angles[1])
-
-    axis_2 = wp.quat_rotate(q_1 * q_0, local_2)
-
-    # convert angular velocity to local space
-    w_err_p = wp.quat_rotate_inv(q_p, w_err)
-
-    # given joint axes and angular velocity error, solve for joint velocities
-    c12 = wp.cross(axis_1, axis_2)
-    c02 = wp.cross(axis_0, axis_2)
-
-    vel = wp.vec2(wp.dot(w_err_p, c12) / wp.dot(axis_0, c12), wp.dot(w_err_p, c02) / wp.dot(axis_1, c02))
-
-    return wp.vec2(angles[0], angles[1]), vel
-
-
-@wp.func
-def compute_3d_rotational_dofs(
-    axis_0: wp.vec3,
-    axis_1: wp.vec3,
-    axis_2: wp.vec3,
-    q0: float,
-    q1: float,
-    q2: float,
-    qd0: float,
-    qd1: float,
-    qd2: float,
-):
-    """
-    Computes the rotation quaternion and 3D angular velocity given the joint axes, coordinates and velocities.
-    """
-    q_off = wp.quat_from_matrix(wp.matrix_from_cols(axis_0, axis_1, axis_2))
-
-    # body local axes
-    local_0 = wp.quat_rotate(q_off, wp.vec3(1.0, 0.0, 0.0))
-    local_1 = wp.quat_rotate(q_off, wp.vec3(0.0, 1.0, 0.0))
-    local_2 = wp.quat_rotate(q_off, wp.vec3(0.0, 0.0, 1.0))
-
-    # reconstruct rotation axes
-    axis_0 = local_0
-    q_0 = wp.quat_from_axis_angle(axis_0, q0)
-
-    axis_1 = wp.quat_rotate(q_0, local_1)
-    q_1 = wp.quat_from_axis_angle(axis_1, q1)
-
-    axis_2 = wp.quat_rotate(q_1 * q_0, local_2)
-    q_2 = wp.quat_from_axis_angle(axis_2, q2)
-
-    rot = q_2 * q_1 * q_0
-    vel = axis_0 * qd0 + axis_1 * qd1 + axis_2 * qd2
-
-    return rot, vel
-
-
-@wp.func
-def invert_3d_rotational_dofs(
-    axis_0: wp.vec3, axis_1: wp.vec3, axis_2: wp.vec3, q_p: wp.quat, q_c: wp.quat, w_err: wp.vec3
-):
-    """
-    Computes generalized joint position and velocity coordinates for a 3D rotational joint given the joint axes, relative orientations and angular velocity differences between the two bodies the joint connects.
-    """
-    q_off = wp.quat_from_matrix(wp.matrix_from_cols(axis_0, axis_1, axis_2))
-    q_pc = wp.quat_inverse(q_off) * wp.quat_inverse(q_p) * q_c * q_off
-
-    # decompose to a compound rotation each axis
-    angles = quat_decompose(q_pc)
-
-    # find rotation axes
-    local_0 = wp.quat_rotate(q_off, wp.vec3(1.0, 0.0, 0.0))
-    local_1 = wp.quat_rotate(q_off, wp.vec3(0.0, 1.0, 0.0))
-    local_2 = wp.quat_rotate(q_off, wp.vec3(0.0, 0.0, 1.0))
-
-    axis_0 = local_0
-    q_0 = wp.quat_from_axis_angle(axis_0, angles[0])
-
-    axis_1 = wp.quat_rotate(q_0, local_1)
-    q_1 = wp.quat_from_axis_angle(axis_1, angles[1])
-
-    axis_2 = wp.quat_rotate(q_1 * q_0, local_2)
-
-    # convert angular velocity to local space
-    w_err_p = wp.quat_rotate_inv(q_p, w_err)
-
-    # given joint axes and angular velocity error, solve for joint velocities
-    c12 = wp.cross(axis_1, axis_2)
-    c02 = wp.cross(axis_0, axis_2)
-    c01 = wp.cross(axis_0, axis_1)
-
-    velocities = wp.vec3(
-        wp.dot(w_err_p, c12) / wp.dot(axis_0, c12),
-        wp.dot(w_err_p, c02) / wp.dot(axis_1, c02),
-        wp.dot(w_err_p, c01) / wp.dot(axis_2, c01),
-    )
-
-    return angles, velocities
-
-
-@wp.func
-def eval_single_articulation_fk(
-    joint_start: int,
-    joint_end: int,
-    joint_q: wp.array(dtype=float),
-    joint_qd: wp.array(dtype=float),
-    joint_q_start: wp.array(dtype=int),
-    joint_qd_start: wp.array(dtype=int),
-    joint_type: wp.array(dtype=int),
-    joint_parent: wp.array(dtype=int),
-    joint_child: wp.array(dtype=int),
-    joint_X_p: wp.array(dtype=wp.transform),
-    joint_X_c: wp.array(dtype=wp.transform),
-    joint_axis: wp.array(dtype=wp.vec3),
-    joint_axis_start: wp.array(dtype=int),
-    joint_axis_dim: wp.array(dtype=int, ndim=2),
-    body_com: wp.array(dtype=wp.vec3),
-    # outputs
-    body_q: wp.array(dtype=wp.transform),
-    body_qd: wp.array(dtype=wp.spatial_vector),
-):
-    for i in range(joint_start, joint_end):
-        parent = joint_parent[i]
-        child = joint_child[i]
-
-        # compute transform across the joint
-        type = joint_type[i]
-
-        X_pj = joint_X_p[i]
-        X_cj = joint_X_c[i]
-
-        # parent anchor frame in world space
-        X_wpj = X_pj
-        # velocity of parent anchor point in world space
-        v_wpj = wp.spatial_vector()
-        if parent >= 0:
-            X_wp = body_q[parent]
-            X_wpj = X_wp * X_wpj
-            r_p = wp.transform_get_translation(X_wpj) - wp.transform_point(X_wp, body_com[parent])
-
-            v_wp = body_qd[parent]
-            w_p = wp.spatial_top(v_wp)
-            v_p = wp.spatial_bottom(v_wp) + wp.cross(w_p, r_p)
-            v_wpj = wp.spatial_vector(w_p, v_p)
-
-        q_start = joint_q_start[i]
-        qd_start = joint_qd_start[i]
-        axis_start = joint_axis_start[i]
-        lin_axis_count = joint_axis_dim[i, 0]
-        ang_axis_count = joint_axis_dim[i, 1]
-
-        X_j = wp.transform_identity()
-        v_j = wp.spatial_vector(wp.vec3(), wp.vec3())
-
-        if type == wp.sim.JOINT_PRISMATIC:
-            axis = joint_axis[axis_start]
-
-            q = joint_q[q_start]
-            qd = joint_qd[qd_start]
-
-            X_j = wp.transform(axis * q, wp.quat_identity())
-            v_j = wp.spatial_vector(wp.vec3(), axis * qd)
-
-        if type == wp.sim.JOINT_REVOLUTE:
-            axis = joint_axis[axis_start]
-
-            q = joint_q[q_start]
-            qd = joint_qd[qd_start]
-
-            X_j = wp.transform(wp.vec3(), wp.quat_from_axis_angle(axis, q))
-            v_j = wp.spatial_vector(axis * qd, wp.vec3())
-
-        if type == wp.sim.JOINT_BALL:
-            r = wp.quat(joint_q[q_start + 0], joint_q[q_start + 1], joint_q[q_start + 2], joint_q[q_start + 3])
-
-            w = wp.vec3(joint_qd[qd_start + 0], joint_qd[qd_start + 1], joint_qd[qd_start + 2])
-
-            X_j = wp.transform(wp.vec3(), r)
-            v_j = wp.spatial_vector(w, wp.vec3())
-
-        if type == wp.sim.JOINT_FREE or type == wp.sim.JOINT_DISTANCE:
-            t = wp.transform(
-                wp.vec3(joint_q[q_start + 0], joint_q[q_start + 1], joint_q[q_start + 2]),
-                wp.quat(joint_q[q_start + 3], joint_q[q_start + 4], joint_q[q_start + 5], joint_q[q_start + 6]),
-            )
-
-            v = wp.spatial_vector(
-                wp.vec3(joint_qd[qd_start + 0], joint_qd[qd_start + 1], joint_qd[qd_start + 2]),
-                wp.vec3(joint_qd[qd_start + 3], joint_qd[qd_start + 4], joint_qd[qd_start + 5]),
-            )
-
-            X_j = t
-            v_j = v
-
-        if type == wp.sim.JOINT_COMPOUND:
-            rot, vel_w = compute_3d_rotational_dofs(
-                joint_axis[axis_start],
-                joint_axis[axis_start + 1],
-                joint_axis[axis_start + 2],
-                joint_q[q_start + 0],
-                joint_q[q_start + 1],
-                joint_q[q_start + 2],
-                joint_qd[qd_start + 0],
-                joint_qd[qd_start + 1],
-                joint_qd[qd_start + 2],
-            )
-
-            t = wp.transform(wp.vec3(0.0, 0.0, 0.0), rot)
-            v = wp.spatial_vector(vel_w, wp.vec3(0.0, 0.0, 0.0))
-
-            X_j = t
-            v_j = v
-
-        if type == wp.sim.JOINT_UNIVERSAL:
-            rot, vel_w = compute_2d_rotational_dofs(
-                joint_axis[axis_start],
-                joint_axis[axis_start + 1],
-                joint_q[q_start + 0],
-                joint_q[q_start + 1],
-                joint_qd[qd_start + 0],
-                joint_qd[qd_start + 1],
-            )
-
-            t = wp.transform(wp.vec3(0.0, 0.0, 0.0), rot)
-            v = wp.spatial_vector(vel_w, wp.vec3(0.0, 0.0, 0.0))
-
-            X_j = t
-            v_j = v
-
-        if type == wp.sim.JOINT_D6:
-            pos = wp.vec3(0.0)
-            rot = wp.quat_identity()
-            vel_v = wp.vec3(0.0)
-            vel_w = wp.vec3(0.0)
-
-            # unroll for loop to ensure joint actions remain differentiable
-            # (since differentiating through a for loop that updates a local variable is not supported)
-
-            if lin_axis_count > 0:
-                axis = joint_axis[axis_start + 0]
-                pos += axis * joint_q[q_start + 0]
-                vel_v += axis * joint_qd[qd_start + 0]
-            if lin_axis_count > 1:
-                axis = joint_axis[axis_start + 1]
-                pos += axis * joint_q[q_start + 1]
-                vel_v += axis * joint_qd[qd_start + 1]
-            if lin_axis_count > 2:
-                axis = joint_axis[axis_start + 2]
-                pos += axis * joint_q[q_start + 2]
-                vel_v += axis * joint_qd[qd_start + 2]
-
-            ia = axis_start + lin_axis_count
-            iq = q_start + lin_axis_count
-            iqd = qd_start + lin_axis_count
-            if ang_axis_count == 1:
-                axis = joint_axis[ia]
-                rot = wp.quat_from_axis_angle(axis, joint_q[iq])
-                vel_w = joint_qd[iqd] * axis
-            if ang_axis_count == 2:
-                rot, vel_w = compute_2d_rotational_dofs(
-                    joint_axis[ia + 0],
-                    joint_axis[ia + 1],
-                    joint_q[iq + 0],
-                    joint_q[iq + 1],
-                    joint_qd[iqd + 0],
-                    joint_qd[iqd + 1],
-                )
-            if ang_axis_count == 3:
-                rot, vel_w = compute_3d_rotational_dofs(
-                    joint_axis[ia + 0],
-                    joint_axis[ia + 1],
-                    joint_axis[ia + 2],
-                    joint_q[iq + 0],
-                    joint_q[iq + 1],
-                    joint_q[iq + 2],
-                    joint_qd[iqd + 0],
-                    joint_qd[iqd + 1],
-                    joint_qd[iqd + 2],
-                )
-
-            X_j = wp.transform(pos, rot)
-            v_j = wp.spatial_vector(vel_w, vel_v)
-
-        # transform from world to joint anchor frame at child body
-        X_wcj = X_wpj * X_j
-        # transform from world to child body frame
-        X_wc = X_wcj * wp.transform_inverse(X_cj)
-
-        # transform velocity across the joint to world space
-        angular_vel = wp.transform_vector(X_wpj, wp.spatial_top(v_j))
-        linear_vel = wp.transform_vector(X_wpj, wp.spatial_bottom(v_j))
-
-        v_wc = v_wpj + wp.spatial_vector(angular_vel, linear_vel)
-
-        body_q[child] = X_wc
-        body_qd[child] = v_wc
-
-
-# implementation where mask is an integer array
-@wp.kernel
-def eval_articulation_fk(
-    articulation_start: wp.array(dtype=int),
-    articulation_mask: wp.array(
-        dtype=int
-    ),  # used to enable / disable FK for an articulation, if None then treat all as enabled
-    joint_q: wp.array(dtype=float),
-    joint_qd: wp.array(dtype=float),
-    joint_q_start: wp.array(dtype=int),
-    joint_qd_start: wp.array(dtype=int),
-    joint_type: wp.array(dtype=int),
-    joint_parent: wp.array(dtype=int),
-    joint_child: wp.array(dtype=int),
-    joint_X_p: wp.array(dtype=wp.transform),
-    joint_X_c: wp.array(dtype=wp.transform),
-    joint_axis: wp.array(dtype=wp.vec3),
-    joint_axis_start: wp.array(dtype=int),
-    joint_axis_dim: wp.array(dtype=int, ndim=2),
-    body_com: wp.array(dtype=wp.vec3),
-    # outputs
-    body_q: wp.array(dtype=wp.transform),
-    body_qd: wp.array(dtype=wp.spatial_vector),
-):
-    tid = wp.tid()
-
-    # early out if disabling FK for this articulation
-    if articulation_mask:
-        if articulation_mask[tid] == 0:
-            return
-
-    joint_start = articulation_start[tid]
-    joint_end = articulation_start[tid + 1]
-
-    eval_single_articulation_fk(
-        joint_start,
-        joint_end,
-        joint_q,
-        joint_qd,
-        joint_q_start,
-        joint_qd_start,
-        joint_type,
-        joint_parent,
-        joint_child,
-        joint_X_p,
-        joint_X_c,
-        joint_axis,
-        joint_axis_start,
-        joint_axis_dim,
-        body_com,
-        # outputs
-        body_q,
-        body_qd,
-    )
-
-
-# overload where mask is a bool array
-@wp.kernel
-def eval_articulation_fk(
-    articulation_start: wp.array(dtype=int),
-    articulation_mask: wp.array(
-        dtype=bool
-    ),  # used to enable / disable FK for an articulation, if None then treat all as enabled
-    joint_q: wp.array(dtype=float),
-    joint_qd: wp.array(dtype=float),
-    joint_q_start: wp.array(dtype=int),
-    joint_qd_start: wp.array(dtype=int),
-    joint_type: wp.array(dtype=int),
-    joint_parent: wp.array(dtype=int),
-    joint_child: wp.array(dtype=int),
-    joint_X_p: wp.array(dtype=wp.transform),
-    joint_X_c: wp.array(dtype=wp.transform),
-    joint_axis: wp.array(dtype=wp.vec3),
-    joint_axis_start: wp.array(dtype=int),
-    joint_axis_dim: wp.array(dtype=int, ndim=2),
-    body_com: wp.array(dtype=wp.vec3),
-    # outputs
-    body_q: wp.array(dtype=wp.transform),
-    body_qd: wp.array(dtype=wp.spatial_vector),
-):
-    tid = wp.tid()
-
-    # early out if disabling FK for this articulation
-    if articulation_mask:
-        if not articulation_mask[tid]:
-            return
-
-    joint_start = articulation_start[tid]
-    joint_end = articulation_start[tid + 1]
-
-    eval_single_articulation_fk(
-        joint_start,
-        joint_end,
-        joint_q,
-        joint_qd,
-        joint_q_start,
-        joint_qd_start,
-        joint_type,
-        joint_parent,
-        joint_child,
-        joint_X_p,
-        joint_X_c,
-        joint_axis,
-        joint_axis_start,
-        joint_axis_dim,
-        body_com,
-        # outputs
-        body_q,
-        body_qd,
-    )
-
-
-# updates state body information based on joint coordinates
-def eval_fk(model, joint_q, joint_qd, mask, state):
-    """
-    Evaluates the model's forward kinematics given the joint coordinates and updates the state's body information (:attr:`State.body_q` and :attr:`State.body_qd`).
-
-    Args:
-        model (Model): The model to evaluate.
-        joint_q (array): Generalized joint position coordinates, shape [joint_coord_count], float
-        joint_qd (array): Generalized joint velocity coordinates, shape [joint_dof_count], float
-        mask (array): The mask to use to enable / disable FK for an articulation. If None then treat all as enabled, shape [articulation_count], int/bool
-        state (State): The state to update.
-    """
-    wp.launch(
-        kernel=eval_articulation_fk,
-        dim=model.articulation_count,
-        inputs=[
-            model.articulation_start,
-            mask,
-            joint_q,
-            joint_qd,
-            model.joint_q_start,
-            model.joint_qd_start,
-            model.joint_type,
-            model.joint_parent,
-            model.joint_child,
-            model.joint_X_p,
-            model.joint_X_c,
-            model.joint_axis,
-            model.joint_axis_start,
-            model.joint_axis_dim,
-            model.body_com,
-        ],
-        outputs=[
-            state.body_q,
-            state.body_qd,
-        ],
-        device=model.device,
-    )
-
-
-@wp.func
-def reconstruct_angular_q_qd(q_pc: wp.quat, w_err: wp.vec3, X_wp: wp.transform, axis: wp.vec3):
-    """
-    Reconstructs the angular joint coordinates and velocities given the relative rotation and angular velocity
-    between a parent and child body.
-
-    Args:
-        q_pc (quat): The relative rotation between the parent and child body.
-        w_err (vec3): The angular velocity between the parent and child body.
-        X_wp (transform): The parent body's transform in world space.
-        axis (vec3): The joint axis in the frame of the parent body.
-
-    Returns:
-        q (float): The joint position coordinate.
-        qd (float): The joint velocity coordinate.
-    """
-    axis_p = wp.transform_vector(X_wp, axis)
-    twist = quat_twist(axis, q_pc)
-    q = wp.acos(twist[3]) * 2.0 * wp.sign(wp.dot(axis, wp.vec3(twist[0], twist[1], twist[2])))
-    qd = wp.dot(w_err, axis_p)
-    return q, qd
-
-
-@wp.kernel
-def eval_articulation_ik(
-    body_q: wp.array(dtype=wp.transform),
-    body_qd: wp.array(dtype=wp.spatial_vector),
-    body_com: wp.array(dtype=wp.vec3),
-    joint_type: wp.array(dtype=int),
-    joint_parent: wp.array(dtype=int),
-    joint_child: wp.array(dtype=int),
-    joint_X_p: wp.array(dtype=wp.transform),
-    joint_X_c: wp.array(dtype=wp.transform),
-    joint_axis: wp.array(dtype=wp.vec3),
-    joint_axis_start: wp.array(dtype=int),
-    joint_axis_dim: wp.array(dtype=int, ndim=2),
-    joint_q_start: wp.array(dtype=int),
-    joint_qd_start: wp.array(dtype=int),
-    joint_q: wp.array(dtype=float),
-    joint_qd: wp.array(dtype=float),
-):
-    tid = wp.tid()
-
-    parent = joint_parent[tid]
-    child = joint_child[tid]
-
-    X_pj = joint_X_p[tid]
-    X_cj = joint_X_c[tid]
-
-    w_p = wp.vec3()
-    v_p = wp.vec3()
-    v_wp = wp.spatial_vector()
-
-    # parent anchor frame in world space
-    X_wpj = X_pj
-    if parent >= 0:
-        X_wp = body_q[parent]
-        X_wpj = X_wp * X_pj
-        r_p = wp.transform_get_translation(X_wpj) - wp.transform_point(X_wp, body_com[parent])
-
-        v_wp = body_qd[parent]
-        w_p = wp.spatial_top(v_wp)
-        v_p = wp.spatial_bottom(v_wp) + wp.cross(w_p, r_p)
-
-    # child transform and moment arm
-    X_wc = body_q[child]
-    X_wcj = X_wc * X_cj
-
-    v_wc = body_qd[child]
-
-    w_c = wp.spatial_top(v_wc)
-    v_c = wp.spatial_bottom(v_wc)
-
-    # joint properties
-    type = joint_type[tid]
-
-    # compute position and orientation differences between anchor frames
-    x_p = wp.transform_get_translation(X_wpj)
-    x_c = wp.transform_get_translation(X_wcj)
-
-    q_p = wp.transform_get_rotation(X_wpj)
-    q_c = wp.transform_get_rotation(X_wcj)
-
-    x_err = x_c - x_p
-    v_err = v_c - v_p
-    w_err = w_c - w_p
-
-    q_start = joint_q_start[tid]
-    qd_start = joint_qd_start[tid]
-    axis_start = joint_axis_start[tid]
-    lin_axis_count = joint_axis_dim[tid, 0]
-    ang_axis_count = joint_axis_dim[tid, 1]
-
-    if type == wp.sim.JOINT_PRISMATIC:
-        axis = joint_axis[axis_start]
-
-        # world space joint axis
-        axis_p = wp.quat_rotate(q_p, axis)
-
-        # evaluate joint coordinates
-        q = wp.dot(x_err, axis_p)
-        qd = wp.dot(v_err, axis_p)
-
-        joint_q[q_start] = q
-        joint_qd[qd_start] = qd
-
-        return
-
-    if type == wp.sim.JOINT_REVOLUTE:
-        axis = joint_axis[axis_start]
-        q_pc = wp.quat_inverse(q_p) * q_c
-
-        q, qd = reconstruct_angular_q_qd(q_pc, w_err, X_wpj, axis)
-
-        joint_q[q_start] = q
-        joint_qd[qd_start] = qd
-
-        return
-
-    if type == wp.sim.JOINT_BALL:
-        q_pc = wp.quat_inverse(q_p) * q_c
-
-        joint_q[q_start + 0] = q_pc[0]
-        joint_q[q_start + 1] = q_pc[1]
-        joint_q[q_start + 2] = q_pc[2]
-        joint_q[q_start + 3] = q_pc[3]
-
-        ang_vel = wp.transform_vector(wp.transform_inverse(X_wpj), w_err)
-        joint_qd[qd_start + 0] = ang_vel[0]
-        joint_qd[qd_start + 1] = ang_vel[1]
-        joint_qd[qd_start + 2] = ang_vel[2]
-
-        return
-
-    if type == wp.sim.JOINT_FIXED:
-        return
-
-    if type == wp.sim.JOINT_FREE or type == wp.sim.JOINT_DISTANCE:
-        q_pc = wp.quat_inverse(q_p) * q_c
-
-        x_err_c = wp.quat_rotate_inv(q_p, x_err)
-        v_err_c = wp.quat_rotate_inv(q_p, v_err)
-        w_err_c = wp.quat_rotate_inv(q_p, w_err)
-
-        joint_q[q_start + 0] = x_err_c[0]
-        joint_q[q_start + 1] = x_err_c[1]
-        joint_q[q_start + 2] = x_err_c[2]
-
-        joint_q[q_start + 3] = q_pc[0]
-        joint_q[q_start + 4] = q_pc[1]
-        joint_q[q_start + 5] = q_pc[2]
-        joint_q[q_start + 6] = q_pc[3]
-
-        joint_qd[qd_start + 0] = w_err_c[0]
-        joint_qd[qd_start + 1] = w_err_c[1]
-        joint_qd[qd_start + 2] = w_err_c[2]
-
-        joint_qd[qd_start + 3] = v_err_c[0]
-        joint_qd[qd_start + 4] = v_err_c[1]
-        joint_qd[qd_start + 5] = v_err_c[2]
-
-        return
-
-    if type == wp.sim.JOINT_COMPOUND:
-        axis_0 = joint_axis[axis_start + 0]
-        axis_1 = joint_axis[axis_start + 1]
-        axis_2 = joint_axis[axis_start + 2]
-        qs, qds = invert_3d_rotational_dofs(axis_0, axis_1, axis_2, q_p, q_c, w_err)
-        joint_q[q_start + 0] = qs[0]
-        joint_q[q_start + 1] = qs[1]
-        joint_q[q_start + 2] = qs[2]
-        joint_qd[qd_start + 0] = qds[0]
-        joint_qd[qd_start + 1] = qds[1]
-        joint_qd[qd_start + 2] = qds[2]
-
-        return
-
-    if type == wp.sim.JOINT_UNIVERSAL:
-        axis_0 = joint_axis[axis_start + 0]
-        axis_1 = joint_axis[axis_start + 1]
-        qs2, qds2 = invert_2d_rotational_dofs(axis_0, axis_1, q_p, q_c, w_err)
-        joint_q[q_start + 0] = qs2[0]
-        joint_q[q_start + 1] = qs2[1]
-        joint_qd[qd_start + 0] = qds2[0]
-        joint_qd[qd_start + 1] = qds2[1]
-
-        return
-
-    if type == wp.sim.JOINT_D6:
-        x_err_c = wp.quat_rotate_inv(q_p, x_err)
-        v_err_c = wp.quat_rotate_inv(q_p, v_err)
-        if lin_axis_count > 0:
-            axis = joint_axis[axis_start + 0]
-            joint_q[q_start + 0] = wp.dot(x_err_c, axis)
-            joint_qd[qd_start + 0] = wp.dot(v_err_c, axis)
-
-        if lin_axis_count > 1:
-            axis = joint_axis[axis_start + 1]
-            joint_q[q_start + 1] = wp.dot(x_err_c, axis)
-            joint_qd[qd_start + 1] = wp.dot(v_err_c, axis)
-
-        if lin_axis_count > 2:
-            axis = joint_axis[axis_start + 2]
-            joint_q[q_start + 2] = wp.dot(x_err_c, axis)
-            joint_qd[qd_start + 2] = wp.dot(v_err_c, axis)
-
-        if ang_axis_count == 1:
-            axis = joint_axis[axis_start]
-            q_pc = wp.quat_inverse(q_p) * q_c
-            q, qd = reconstruct_angular_q_qd(q_pc, w_err, X_wpj, joint_axis[axis_start + lin_axis_count])
-            joint_q[q_start + lin_axis_count] = q
-            joint_qd[qd_start + lin_axis_count] = qd
-
-        if ang_axis_count == 2:
-            axis_0 = joint_axis[axis_start + lin_axis_count + 0]
-            axis_1 = joint_axis[axis_start + lin_axis_count + 1]
-            qs2, qds2 = invert_2d_rotational_dofs(axis_0, axis_1, q_p, q_c, w_err)
-            joint_q[q_start + lin_axis_count + 0] = qs2[0]
-            joint_q[q_start + lin_axis_count + 1] = qs2[1]
-            joint_qd[qd_start + lin_axis_count + 0] = qds2[0]
-            joint_qd[qd_start + lin_axis_count + 1] = qds2[1]
-
-        if ang_axis_count == 3:
-            axis_0 = joint_axis[axis_start + lin_axis_count + 0]
-            axis_1 = joint_axis[axis_start + lin_axis_count + 1]
-            axis_2 = joint_axis[axis_start + lin_axis_count + 2]
-            qs3, qds3 = invert_3d_rotational_dofs(axis_0, axis_1, axis_2, q_p, q_c, w_err)
-            joint_q[q_start + lin_axis_count + 0] = qs3[0]
-            joint_q[q_start + lin_axis_count + 1] = qs3[1]
-            joint_q[q_start + lin_axis_count + 2] = qs3[2]
-            joint_qd[qd_start + lin_axis_count + 0] = qds3[0]
-            joint_qd[qd_start + lin_axis_count + 1] = qds3[1]
-            joint_qd[qd_start + lin_axis_count + 2] = qds3[2]
-
-        return
-
-
-# given maximal coordinate model computes ik (closest point projection)
-def eval_ik(model, state, joint_q, joint_qd):
-    """
-    Evaluates the model's inverse kinematics given the state's body information (:attr:`State.body_q` and :attr:`State.body_qd`) and updates the generalized joint coordinates `joint_q` and `joint_qd`.
-
-    Args:
-        model (Model): The model to evaluate.
-        state (State): The state with the body's maximal coordinates (positions :attr:`State.body_q` and velocities :attr:`State.body_qd`) to use.
-        joint_q (array): Generalized joint position coordinates, shape [joint_coord_count], float
-        joint_qd (array): Generalized joint velocity coordinates, shape [joint_dof_count], float
-    """
-    wp.launch(
-        kernel=eval_articulation_ik,
-        dim=model.joint_count,
-        inputs=[
-            state.body_q,
-            state.body_qd,
-            model.body_com,
-            model.joint_type,
-            model.joint_parent,
-            model.joint_child,
-            model.joint_X_p,
-            model.joint_X_c,
-            model.joint_axis,
-            model.joint_axis_start,
-            model.joint_axis_dim,
-            model.joint_q_start,
-            model.joint_qd_start,
-        ],
-        outputs=[joint_q, joint_qd],
-        device=model.device,
-    )