warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/sim/articulation.py

@@ -7,7 +7,161 @@
 
 import warp as wp
 
-from .utils import quat_decompose, quat_twist, quat_twist_angle
+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.mat33(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.mat33(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.mat33(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, todo: can probably use fact that rz'*ry'*rx' == rx*ry*rz to avoid some work here
+    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.mat33(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.kernel
@@ -46,12 +200,6 @@ def eval_articulation_fk(
     for i in range(joint_start, joint_end):
         parent = joint_parent[i]
         child = joint_child[i]
-        X_wp = wp.transform_identity()
-        v_wp = wp.spatial_vector()
-
-        if parent >= 0:
-            X_wp = body_q[parent]
-            v_wp = body_qd[parent]
 
         # compute transform across the joint
         type = joint_type[i]
@@ -59,14 +207,28 @@ def eval_articulation_fk(
         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_jc = wp.transform_identity()
-        v_jc = wp.spatial_vector(wp.vec3(), wp.vec3())
+        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]
@@ -74,8 +236,8 @@ def eval_articulation_fk(
             q = joint_q[q_start]
             qd = joint_qd[qd_start]
 
-            X_jc = wp.transform(axis * q, wp.quat_identity())
-            v_jc = wp.spatial_vector(wp.vec3(), axis * qd)
+            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]
@@ -83,20 +245,18 @@ def eval_articulation_fk(
             q = joint_q[q_start]
             qd = joint_qd[qd_start]
 
-            X_jc = wp.transform(wp.vec3(), wp.quat_from_axis_angle(axis, q))
-            v_jc = wp.spatial_vector(axis * qd, wp.vec3())
+            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])
 
-            # print(r)
-
-            X_jc = wp.transform(wp.vec3(), r)
-            v_jc = wp.spatial_vector(w, wp.vec3())
+            X_j = wp.transform(wp.vec3(), r)
+            v_j = wp.spatial_vector(w, wp.vec3())
 
-        if type == wp.sim.JOINT_FREE:
+        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]),
@@ -107,46 +267,43 @@ def eval_articulation_fk(
                 wp.vec3(joint_qd[qd_start + 3], joint_qd[qd_start + 4], joint_qd[qd_start + 5]),
             )
 
-            X_jc = t
-            v_jc = v
+            X_j = t
+            v_j = v
 
         if type == wp.sim.JOINT_COMPOUND:
-            axis = joint_axis[axis_start]
-
-            # reconstruct rotation axes, todo: can probably use fact that rz'*ry'*rx' == rx*ry*rz to avoid some work here
-            axis_0 = axis
-            q_0 = wp.quat_from_axis_angle(axis_0, joint_q[q_start + 0])
-
-            axis_1 = joint_axis[axis_start + 1]
-            q_1 = wp.quat_from_axis_angle(axis_1, joint_q[q_start + 1])
-
-            axis_2 = joint_axis[axis_start + 2]
-            q_2 = wp.quat_from_axis_angle(axis_2, joint_q[q_start + 2])
-
-            t = wp.transform(wp.vec3(), q_2 * q_1 * q_0)
-
-            v = wp.spatial_vector(
-                axis_0 * joint_qd[qd_start + 0] + axis_1 * joint_qd[qd_start + 1] + axis_2 * joint_qd[qd_start + 2],
-                wp.vec3(),
+            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],
             )
 
-            X_jc = t
-            v_jc = v
+            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))
 
-        if type == wp.sim.JOINT_UNIVERSAL:
-            # reconstruct rotation axes
-            axis_0 = joint_axis[axis_start]
-            q_0 = wp.quat_from_axis_angle(axis_0, joint_q[q_start + 0])
-
-            axis_1 = joint_axis[axis_start + 1]
-            q_1 = wp.quat_from_axis_angle(axis_1, joint_q[q_start + 1])
+            X_j = t
+            v_j = v
 
-            t = wp.transform(wp.vec3(), q_1 * q_0)
+        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],
+            )
 
-            v = wp.spatial_vector(axis_0 * joint_qd[qd_start + 0] + axis_1 * joint_qd[qd_start + 1], wp.vec3())
+            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_jc = t
-            v_jc = v
+            X_j = t
+            v_j = v
 
         if type == wp.sim.JOINT_D6:
             pos = wp.vec3(0.0)
@@ -170,34 +327,49 @@ def eval_articulation_fk(
                 pos += axis * joint_q[q_start + 2]
                 vel_v += axis * joint_qd[qd_start + 2]
 
-            if ang_axis_count > 0:
+            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[axis_start + lin_axis_count + 0]
                 qi = wp.quat_from_axis_angle(axis, joint_q[q_start + lin_axis_count + 0])
                 rot = qi * rot
-                vel_w += joint_qd[qd_start + lin_axis_count + 0] * axis
-            if ang_axis_count > 1:
-                axis = joint_axis[axis_start + lin_axis_count + 1]
-                qi = wp.quat_from_axis_angle(axis, joint_q[q_start + lin_axis_count + 1])
-                rot = qi * rot
-                vel_w += joint_qd[qd_start + lin_axis_count + 1] * axis
-            if ang_axis_count > 2:
-                axis = joint_axis[axis_start + lin_axis_count + 2]
-                qi = wp.quat_from_axis_angle(axis, joint_q[q_start + lin_axis_count + 2])
-                rot = qi * rot
-                vel_w += joint_qd[qd_start + lin_axis_count + 2] * axis
-
-            X_jc = wp.transform(pos, rot)
-            v_jc = wp.spatial_vector(vel_w, vel_v)
-
-        X_wj = X_wp * X_pj
-        X_wc = X_wj * X_jc
-        X_wc *= wp.transform_inverse(X_cj)
+                vel_w = joint_qd[qd_start + lin_axis_count + 0] * 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_wj, wp.spatial_top(v_jc))
-        linear_vel = wp.transform_vector(X_wj, wp.spatial_bottom(v_jc))
+        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_wp + wp.spatial_vector(angular_vel, linear_vel + wp.cross(angular_vel, body_com[i]))
+        v_wc = v_wpj + wp.spatial_vector(angular_vel, linear_vel)
 
         body_q[child] = X_wc
         body_qd[child] = v_wc
@@ -205,6 +377,16 @@ def eval_articulation_fk(
 
 # 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
+        state (State): The state to update.
+    """
     wp.launch(
         kernel=eval_articulation_fk,
         dim=model.articulation_count,
@@ -233,6 +415,29 @@ def eval_fk(model, joint_q, joint_qd, mask, state):
     )
 
 
+@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),
@@ -253,46 +458,46 @@ def eval_articulation_ik(
 ):
     tid = wp.tid()
 
-    c_parent = joint_parent[tid]
-    c_child = joint_child[tid]
+    parent = joint_parent[tid]
+    child = joint_child[tid]
 
     X_pj = joint_X_p[tid]
     X_cj = joint_X_c[tid]
 
-    X_wp = X_pj
-    r_p = wp.vec3()
     w_p = wp.vec3()
     v_p = wp.vec3()
+    v_wp = wp.spatial_vector()
 
-    # parent transform and moment arm
-    if c_parent >= 0:
-        X_wp = body_q[c_parent] * X_wp
-        r_wp = wp.transform_get_translation(X_wp) - wp.transform_point(body_q[c_parent], body_com[c_parent])
+    # 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])
 
-        twist_p = body_qd[c_parent]
-
-        w_p = wp.spatial_top(twist_p)
-        v_p = wp.spatial_bottom(twist_p) + wp.cross(w_p, r_wp)
+        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[c_child] * joint_X_c[tid]
-    r_c = wp.transform_get_translation(X_wc) - wp.transform_point(body_q[c_child], body_com[c_child])
+    X_wc = body_q[child]
+    X_wcj = X_wc * X_cj
 
-    twist_c = body_qd[c_child]
+    v_wc = body_qd[child]
 
-    w_c = wp.spatial_top(twist_c)
-    v_c = wp.spatial_bottom(twist_c) + wp.cross(w_c, r_c)
+    w_c = wp.spatial_top(v_wc)
+    v_c = wp.spatial_bottom(v_wc)
 
     # joint properties
     type = joint_type[tid]
 
-    x_p = wp.transform_get_translation(X_wp)
-    x_c = wp.transform_get_translation(X_wc)
+    # 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_wp)
-    q_c = wp.transform_get_rotation(X_wc)
+    q_p = wp.transform_get_rotation(X_wpj)
+    q_c = wp.transform_get_rotation(X_wcj)
 
-    # translational error
     x_err = x_c - x_p
     v_err = v_c - v_p
     w_err = w_c - w_p
@@ -307,7 +512,7 @@ def eval_articulation_ik(
         axis = joint_axis[axis_start]
 
         # world space joint axis
-        axis_p = wp.transform_vector(X_wp, axis)
+        axis_p = wp.quat_rotate(q_p, axis)
 
         # evaluate joint coordinates
         q = wp.dot(x_err, axis_p)
@@ -320,16 +525,9 @@ def eval_articulation_ik(
 
     if type == wp.sim.JOINT_REVOLUTE:
         axis = joint_axis[axis_start]
-
-        axis_p = wp.transform_vector(X_wp, axis)
-        axis_c = wp.transform_vector(X_wc, axis)
-
-        # swing twist decomposition
         q_pc = wp.quat_inverse(q_p) * q_c
-        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)
+        q, qd = reconstruct_angular_q_qd(q_pc, w_err, X_wpj, axis)
 
         joint_q[q_start] = q
         joint_qd[qd_start] = qd
@@ -344,125 +542,127 @@ def eval_articulation_ik(
         joint_q[q_start + 2] = q_pc[2]
         joint_q[q_start + 3] = q_pc[3]
 
-        joint_qd[qd_start + 0] = w_err[0]
-        joint_qd[qd_start + 1] = w_err[1]
-        joint_qd[qd_start + 2] = w_err[2]
+        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:
+    if type == wp.sim.JOINT_FREE or type == wp.sim.JOINT_DISTANCE:
         q_pc = wp.quat_inverse(q_p) * q_c
 
-        joint_q[q_start + 0] = x_err[0]
-        joint_q[q_start + 1] = x_err[1]
-        joint_q[q_start + 2] = x_err[2]
+        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[0]
-        joint_qd[qd_start + 1] = w_err[1]
-        joint_qd[qd_start + 2] = w_err[2]
-
-        joint_qd[qd_start + 3] = v_err[0]
-        joint_qd[qd_start + 4] = v_err[1]
-        joint_qd[qd_start + 5] = v_err[2]
-
-    if type == wp.sim.JOINT_COMPOUND:
-        q_off = wp.transform_get_rotation(X_cj)
-        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)
-
-        # reconstruct rotation axes
-        axis_0 = wp.vec3(1.0, 0.0, 0.0)
-        q_0 = wp.quat_from_axis_angle(axis_0, angles[0])
+        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]
 
-        axis_1 = wp.quat_rotate(q_0, wp.vec3(0.0, 1.0, 0.0))
-        q_1 = wp.quat_from_axis_angle(axis_1, angles[1])
+        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]
 
-        axis_2 = wp.quat_rotate(q_1 * q_0, wp.vec3(0.0, 0.0, 1.0))
-        q_2 = wp.quat_from_axis_angle(axis_2, angles[2])
-
-        q_w = q_p * q_off
-
-        joint_q[q_start + 0] = angles[0]
-        joint_q[q_start + 1] = angles[1]
-        joint_q[q_start + 2] = angles[2]
+        return
 
-        joint_qd[qd_start + 0] = wp.dot(wp.quat_rotate(q_w, axis_0), w_err)
-        joint_qd[qd_start + 1] = wp.dot(wp.quat_rotate(q_w, axis_1), w_err)
-        joint_qd[qd_start + 2] = wp.dot(wp.quat_rotate(q_w, axis_2), w_err)
+    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:
-        q_off = wp.transform_get_rotation(X_cj)
-        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)
-
-        # reconstruct rotation axes
-        axis_0 = wp.vec3(1.0, 0.0, 0.0)
-        q_0 = wp.quat_from_axis_angle(axis_0, angles[0])
-
-        axis_1 = wp.quat_rotate(q_0, wp.vec3(0.0, 1.0, 0.0))
-        q_1 = wp.quat_from_axis_angle(axis_1, angles[1])
-
-        q_w = q_p * q_off
-
-        joint_q[q_start + 0] = angles[0]
-        joint_q[q_start + 1] = angles[1]
-
-        joint_qd[qd_start + 0] = wp.dot(wp.quat_rotate(q_w, axis_0), w_err)
-        joint_qd[qd_start + 1] = wp.dot(wp.quat_rotate(q_w, axis_1), w_err)
+        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 = wp.transform_vector(X_wp, joint_axis[axis_start + 0])
-            joint_q[q_start + 0] = wp.dot(x_err, axis)
-            joint_qd[qd_start + 0] = wp.dot(v_err, axis)
+            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 = wp.transform_vector(X_wp, joint_axis[axis_start + 1])
-            joint_q[q_start + 1] = wp.dot(x_err, axis)
-            joint_qd[qd_start + 1] = wp.dot(v_err, axis)
+            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 = wp.transform_vector(X_wp, joint_axis[axis_start + 2])
-            joint_q[q_start + 2] = wp.dot(x_err, axis)
-            joint_qd[qd_start + 2] = wp.dot(v_err, axis)
-
-        q_pc = wp.quat_inverse(q_p) * q_c
-        if ang_axis_count > 0:
-            axis = wp.transform_vector(X_wp, joint_axis[axis_start + lin_axis_count + 0])
-            joint_q[q_start + lin_axis_count + 0] = quat_twist_angle(axis, q_pc)
-            joint_qd[qd_start + lin_axis_count + 0] = wp.dot(w_err, axis)
+            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 = wp.transform_vector(X_wp, joint_axis[axis_start + lin_axis_count + 1])
-            joint_q[q_start + lin_axis_count + 1] = quat_twist_angle(axis, q_pc)
-            joint_qd[qd_start + lin_axis_count + 1] = wp.dot(w_err, axis)
-
-        if ang_axis_count > 2:
-            axis = wp.transform_vector(X_wp, joint_axis[axis_start + lin_axis_count + 2])
-            joint_q[q_start + lin_axis_count + 2] = quat_twist_angle(axis, q_pc)
-            joint_qd[qd_start + lin_axis_count + 2] = wp.dot(w_err, 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,