warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/sim/integrator_featherstone.py

@@ -1,2101 +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 .articulation import (
-    compute_2d_rotational_dofs,
-    compute_3d_rotational_dofs,
-    eval_fk,
-)
-from .integrator import Integrator
-from .integrator_euler import (
-    eval_bending_forces,
-    eval_joint_force,
-    eval_muscle_forces,
-    eval_particle_body_contact_forces,
-    eval_particle_forces,
-    eval_particle_ground_contact_forces,
-    eval_rigid_contacts,
-    eval_spring_forces,
-    eval_tetrahedral_forces,
-    eval_triangle_contact_forces,
-    eval_triangle_forces,
-)
-from .model import Control, Model, State
-
-
-# Frank & Park definition 3.20, pg 100
-@wp.func
-def transform_twist(t: wp.transform, x: wp.spatial_vector):
-    q = wp.transform_get_rotation(t)
-    p = wp.transform_get_translation(t)
-
-    w = wp.spatial_top(x)
-    v = wp.spatial_bottom(x)
-
-    w = wp.quat_rotate(q, w)
-    v = wp.quat_rotate(q, v) + wp.cross(p, w)
-
-    return wp.spatial_vector(w, v)
-
-
-@wp.func
-def transform_wrench(t: wp.transform, x: wp.spatial_vector):
-    q = wp.transform_get_rotation(t)
-    p = wp.transform_get_translation(t)
-
-    w = wp.spatial_top(x)
-    v = wp.spatial_bottom(x)
-
-    v = wp.quat_rotate(q, v)
-    w = wp.quat_rotate(q, w) + wp.cross(p, v)
-
-    return wp.spatial_vector(w, v)
-
-
-@wp.func
-def spatial_adjoint(R: wp.mat33, S: wp.mat33):
-    # T = [R  0]
-    #     [S  R]
-
-    # fmt: off
-    return wp.spatial_matrix(
-        R[0, 0], R[0, 1], R[0, 2],     0.0,     0.0,     0.0,
-        R[1, 0], R[1, 1], R[1, 2],     0.0,     0.0,     0.0,
-        R[2, 0], R[2, 1], R[2, 2],     0.0,     0.0,     0.0,
-        S[0, 0], S[0, 1], S[0, 2], R[0, 0], R[0, 1], R[0, 2],
-        S[1, 0], S[1, 1], S[1, 2], R[1, 0], R[1, 1], R[1, 2],
-        S[2, 0], S[2, 1], S[2, 2], R[2, 0], R[2, 1], R[2, 2],
-    )
-    # fmt: on
-
-
-@wp.kernel
-def compute_spatial_inertia(
-    body_inertia: wp.array(dtype=wp.mat33),
-    body_mass: wp.array(dtype=float),
-    # outputs
-    body_I_m: wp.array(dtype=wp.spatial_matrix),
-):
-    tid = wp.tid()
-    I = body_inertia[tid]
-    m = body_mass[tid]
-    # fmt: off
-    body_I_m[tid] = wp.spatial_matrix(
-        I[0, 0], I[0, 1], I[0, 2], 0.0, 0.0, 0.0,
-        I[1, 0], I[1, 1], I[1, 2], 0.0, 0.0, 0.0,
-        I[2, 0], I[2, 1], I[2, 2], 0.0, 0.0, 0.0,
-        0.0,     0.0,     0.0,     m,   0.0, 0.0,
-        0.0,     0.0,     0.0,     0.0, m,   0.0,
-        0.0,     0.0,     0.0,     0.0, 0.0, m,
-    )
-    # fmt: on
-
-
-@wp.kernel
-def compute_com_transforms(
-    body_com: wp.array(dtype=wp.vec3),
-    # outputs
-    body_X_com: wp.array(dtype=wp.transform),
-):
-    tid = wp.tid()
-    com = body_com[tid]
-    body_X_com[tid] = wp.transform(com, wp.quat_identity())
-
-
-# computes adj_t^-T*I*adj_t^-1 (tensor change of coordinates), Frank & Park, section 8.2.3, pg 290
-@wp.func
-def spatial_transform_inertia(t: wp.transform, I: wp.spatial_matrix):
-    t_inv = wp.transform_inverse(t)
-
-    q = wp.transform_get_rotation(t_inv)
-    p = wp.transform_get_translation(t_inv)
-
-    r1 = wp.quat_rotate(q, wp.vec3(1.0, 0.0, 0.0))
-    r2 = wp.quat_rotate(q, wp.vec3(0.0, 1.0, 0.0))
-    r3 = wp.quat_rotate(q, wp.vec3(0.0, 0.0, 1.0))
-
-    R = wp.matrix_from_cols(r1, r2, r3)
-    S = wp.skew(p) @ R
-
-    T = spatial_adjoint(R, S)
-
-    return wp.mul(wp.mul(wp.transpose(T), I), T)
-
-
-# compute transform across a joint
-@wp.func
-def jcalc_transform(
-    type: int,
-    joint_axis: wp.array(dtype=wp.vec3),
-    axis_start: int,
-    lin_axis_count: int,
-    ang_axis_count: int,
-    joint_q: wp.array(dtype=float),
-    start: int,
-):
-    if type == wp.sim.JOINT_PRISMATIC:
-        q = joint_q[start]
-        axis = joint_axis[axis_start]
-        X_jc = wp.transform(axis * q, wp.quat_identity())
-        return X_jc
-
-    if type == wp.sim.JOINT_REVOLUTE:
-        q = joint_q[start]
-        axis = joint_axis[axis_start]
-        X_jc = wp.transform(wp.vec3(), wp.quat_from_axis_angle(axis, q))
-        return X_jc
-
-    if type == wp.sim.JOINT_BALL:
-        qx = joint_q[start + 0]
-        qy = joint_q[start + 1]
-        qz = joint_q[start + 2]
-        qw = joint_q[start + 3]
-
-        X_jc = wp.transform(wp.vec3(), wp.quat(qx, qy, qz, qw))
-        return X_jc
-
-    if type == wp.sim.JOINT_FIXED:
-        X_jc = wp.transform_identity()
-        return X_jc
-
-    if type == wp.sim.JOINT_FREE or type == wp.sim.JOINT_DISTANCE:
-        px = joint_q[start + 0]
-        py = joint_q[start + 1]
-        pz = joint_q[start + 2]
-
-        qx = joint_q[start + 3]
-        qy = joint_q[start + 4]
-        qz = joint_q[start + 5]
-        qw = joint_q[start + 6]
-
-        X_jc = wp.transform(wp.vec3(px, py, pz), wp.quat(qx, qy, qz, qw))
-        return X_jc
-
-    if type == wp.sim.JOINT_COMPOUND:
-        rot, _ = compute_3d_rotational_dofs(
-            joint_axis[axis_start],
-            joint_axis[axis_start + 1],
-            joint_axis[axis_start + 2],
-            joint_q[start + 0],
-            joint_q[start + 1],
-            joint_q[start + 2],
-            0.0,
-            0.0,
-            0.0,
-        )
-
-        X_jc = wp.transform(wp.vec3(), rot)
-        return X_jc
-
-    if type == wp.sim.JOINT_UNIVERSAL:
-        rot, _ = compute_2d_rotational_dofs(
-            joint_axis[axis_start],
-            joint_axis[axis_start + 1],
-            joint_q[start + 0],
-            joint_q[start + 1],
-            0.0,
-            0.0,
-        )
-
-        X_jc = wp.transform(wp.vec3(), rot)
-        return X_jc
-
-    if type == wp.sim.JOINT_D6:
-        pos = wp.vec3(0.0)
-        rot = wp.quat_identity()
-
-        # 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[start + 0]
-        if lin_axis_count > 1:
-            axis = joint_axis[axis_start + 1]
-            pos += axis * joint_q[start + 1]
-        if lin_axis_count > 2:
-            axis = joint_axis[axis_start + 2]
-            pos += axis * joint_q[start + 2]
-
-        ia = axis_start + lin_axis_count
-        iq = start + lin_axis_count
-        if ang_axis_count == 1:
-            axis = joint_axis[ia]
-            rot = wp.quat_from_axis_angle(axis, joint_q[iq])
-        if ang_axis_count == 2:
-            rot, _ = compute_2d_rotational_dofs(
-                joint_axis[ia + 0],
-                joint_axis[ia + 1],
-                joint_q[iq + 0],
-                joint_q[iq + 1],
-                0.0,
-                0.0,
-            )
-        if ang_axis_count == 3:
-            rot, _ = 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],
-                0.0,
-                0.0,
-                0.0,
-            )
-
-        X_jc = wp.transform(pos, rot)
-        return X_jc
-
-    # default case
-    return wp.transform_identity()
-
-
-# compute motion subspace and velocity for a joint
-@wp.func
-def jcalc_motion(
-    type: int,
-    joint_axis: wp.array(dtype=wp.vec3),
-    axis_start: int,
-    lin_axis_count: int,
-    ang_axis_count: int,
-    X_sc: wp.transform,
-    joint_q: wp.array(dtype=float),
-    joint_qd: wp.array(dtype=float),
-    q_start: int,
-    qd_start: int,
-    # outputs
-    joint_S_s: wp.array(dtype=wp.spatial_vector),
-):
-    if type == wp.sim.JOINT_PRISMATIC:
-        axis = joint_axis[axis_start]
-        S_s = transform_twist(X_sc, wp.spatial_vector(wp.vec3(), axis))
-        v_j_s = S_s * joint_qd[qd_start]
-        joint_S_s[qd_start] = S_s
-        return v_j_s
-
-    if type == wp.sim.JOINT_REVOLUTE:
-        axis = joint_axis[axis_start]
-        S_s = transform_twist(X_sc, wp.spatial_vector(axis, wp.vec3()))
-        v_j_s = S_s * joint_qd[qd_start]
-        joint_S_s[qd_start] = S_s
-        return v_j_s
-
-    if type == wp.sim.JOINT_UNIVERSAL:
-        axis_0 = joint_axis[axis_start + 0]
-        axis_1 = joint_axis[axis_start + 1]
-        q_off = wp.quat_from_matrix(wp.matrix_from_cols(axis_0, axis_1, wp.cross(axis_0, axis_1)))
-        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, joint_q[q_start + 0])
-
-        axis_1 = wp.quat_rotate(q_0, local_1)
-
-        S_0 = transform_twist(X_sc, wp.spatial_vector(axis_0, wp.vec3()))
-        S_1 = transform_twist(X_sc, wp.spatial_vector(axis_1, wp.vec3()))
-
-        joint_S_s[qd_start + 0] = S_0
-        joint_S_s[qd_start + 1] = S_1
-
-        return S_0 * joint_qd[qd_start + 0] + S_1 * joint_qd[qd_start + 1]
-
-    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]
-        q_off = wp.quat_from_matrix(wp.matrix_from_cols(axis_0, axis_1, axis_2))
-        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, joint_q[q_start + 0])
-
-        axis_1 = wp.quat_rotate(q_0, local_1)
-        q_1 = wp.quat_from_axis_angle(axis_1, joint_q[q_start + 1])
-
-        axis_2 = wp.quat_rotate(q_1 * q_0, local_2)
-
-        S_0 = transform_twist(X_sc, wp.spatial_vector(axis_0, wp.vec3()))
-        S_1 = transform_twist(X_sc, wp.spatial_vector(axis_1, wp.vec3()))
-        S_2 = transform_twist(X_sc, wp.spatial_vector(axis_2, wp.vec3()))
-
-        joint_S_s[qd_start + 0] = S_0
-        joint_S_s[qd_start + 1] = S_1
-        joint_S_s[qd_start + 2] = S_2
-
-        return S_0 * joint_qd[qd_start + 0] + S_1 * joint_qd[qd_start + 1] + S_2 * joint_qd[qd_start + 2]
-
-    if type == wp.sim.JOINT_D6:
-        v_j_s = wp.spatial_vector()
-        if lin_axis_count > 0:
-            axis = joint_axis[axis_start + 0]
-            S_s = transform_twist(X_sc, wp.spatial_vector(wp.vec3(), axis))
-            v_j_s += S_s * joint_qd[qd_start + 0]
-            joint_S_s[qd_start + 0] = S_s
-        if lin_axis_count > 1:
-            axis = joint_axis[axis_start + 1]
-            S_s = transform_twist(X_sc, wp.spatial_vector(wp.vec3(), axis))
-            v_j_s += S_s * joint_qd[qd_start + 1]
-            joint_S_s[qd_start + 1] = S_s
-        if lin_axis_count > 2:
-            axis = joint_axis[axis_start + 2]
-            S_s = transform_twist(X_sc, wp.spatial_vector(wp.vec3(), axis))
-            v_j_s += S_s * joint_qd[qd_start + 2]
-            joint_S_s[qd_start + 2] = S_s
-        if ang_axis_count > 0:
-            axis = joint_axis[axis_start + lin_axis_count + 0]
-            S_s = transform_twist(X_sc, wp.spatial_vector(axis, wp.vec3()))
-            v_j_s += S_s * joint_qd[qd_start + lin_axis_count + 0]
-            joint_S_s[qd_start + lin_axis_count + 0] = S_s
-        if ang_axis_count > 1:
-            axis = joint_axis[axis_start + lin_axis_count + 1]
-            S_s = transform_twist(X_sc, wp.spatial_vector(axis, wp.vec3()))
-            v_j_s += S_s * joint_qd[qd_start + lin_axis_count + 1]
-            joint_S_s[qd_start + lin_axis_count + 1] = S_s
-        if ang_axis_count > 2:
-            axis = joint_axis[axis_start + lin_axis_count + 2]
-            S_s = transform_twist(X_sc, wp.spatial_vector(axis, wp.vec3()))
-            v_j_s += S_s * joint_qd[qd_start + lin_axis_count + 2]
-            joint_S_s[qd_start + lin_axis_count + 2] = S_s
-
-        return v_j_s
-
-    if type == wp.sim.JOINT_BALL:
-        S_0 = transform_twist(X_sc, wp.spatial_vector(1.0, 0.0, 0.0, 0.0, 0.0, 0.0))
-        S_1 = transform_twist(X_sc, wp.spatial_vector(0.0, 1.0, 0.0, 0.0, 0.0, 0.0))
-        S_2 = transform_twist(X_sc, wp.spatial_vector(0.0, 0.0, 1.0, 0.0, 0.0, 0.0))
-
-        joint_S_s[qd_start + 0] = S_0
-        joint_S_s[qd_start + 1] = S_1
-        joint_S_s[qd_start + 2] = S_2
-
-        return S_0 * joint_qd[qd_start + 0] + S_1 * joint_qd[qd_start + 1] + S_2 * joint_qd[qd_start + 2]
-
-    if type == wp.sim.JOINT_FIXED:
-        return wp.spatial_vector()
-
-    if type == wp.sim.JOINT_FREE or type == wp.sim.JOINT_DISTANCE:
-        v_j_s = transform_twist(
-            X_sc,
-            wp.spatial_vector(
-                joint_qd[qd_start + 0],
-                joint_qd[qd_start + 1],
-                joint_qd[qd_start + 2],
-                joint_qd[qd_start + 3],
-                joint_qd[qd_start + 4],
-                joint_qd[qd_start + 5],
-            ),
-        )
-
-        joint_S_s[qd_start + 0] = transform_twist(X_sc, wp.spatial_vector(1.0, 0.0, 0.0, 0.0, 0.0, 0.0))
-        joint_S_s[qd_start + 1] = transform_twist(X_sc, wp.spatial_vector(0.0, 1.0, 0.0, 0.0, 0.0, 0.0))
-        joint_S_s[qd_start + 2] = transform_twist(X_sc, wp.spatial_vector(0.0, 0.0, 1.0, 0.0, 0.0, 0.0))
-        joint_S_s[qd_start + 3] = transform_twist(X_sc, wp.spatial_vector(0.0, 0.0, 0.0, 1.0, 0.0, 0.0))
-        joint_S_s[qd_start + 4] = transform_twist(X_sc, wp.spatial_vector(0.0, 0.0, 0.0, 0.0, 1.0, 0.0))
-        joint_S_s[qd_start + 5] = transform_twist(X_sc, wp.spatial_vector(0.0, 0.0, 0.0, 0.0, 0.0, 1.0))
-
-        return v_j_s
-
-    wp.printf("jcalc_motion not implemented for joint type %d\n", type)
-
-    # default case
-    return wp.spatial_vector()
-
-
-# computes joint space forces/torques in tau
-@wp.func
-def jcalc_tau(
-    type: int,
-    joint_target_ke: wp.array(dtype=float),
-    joint_target_kd: wp.array(dtype=float),
-    joint_limit_ke: wp.array(dtype=float),
-    joint_limit_kd: wp.array(dtype=float),
-    joint_S_s: wp.array(dtype=wp.spatial_vector),
-    joint_q: wp.array(dtype=float),
-    joint_qd: wp.array(dtype=float),
-    joint_act: wp.array(dtype=float),
-    joint_axis_mode: wp.array(dtype=int),
-    joint_limit_lower: wp.array(dtype=float),
-    joint_limit_upper: wp.array(dtype=float),
-    coord_start: int,
-    dof_start: int,
-    axis_start: int,
-    lin_axis_count: int,
-    ang_axis_count: int,
-    body_f_s: wp.spatial_vector,
-    # outputs
-    tau: wp.array(dtype=float),
-):
-    if type == wp.sim.JOINT_PRISMATIC or type == wp.sim.JOINT_REVOLUTE:
-        S_s = joint_S_s[dof_start]
-
-        q = joint_q[coord_start]
-        qd = joint_qd[dof_start]
-        act = joint_act[axis_start]
-
-        lower = joint_limit_lower[axis_start]
-        upper = joint_limit_upper[axis_start]
-
-        limit_ke = joint_limit_ke[axis_start]
-        limit_kd = joint_limit_kd[axis_start]
-        target_ke = joint_target_ke[axis_start]
-        target_kd = joint_target_kd[axis_start]
-        mode = joint_axis_mode[axis_start]
-
-        # total torque / force on the joint
-        t = -wp.dot(S_s, body_f_s) + eval_joint_force(
-            q, qd, act, target_ke, target_kd, lower, upper, limit_ke, limit_kd, mode
-        )
-
-        tau[dof_start] = t
-
-        return
-
-    if type == wp.sim.JOINT_BALL:
-        # target_ke = joint_target_ke[axis_start]
-        # target_kd = joint_target_kd[axis_start]
-
-        for i in range(3):
-            S_s = joint_S_s[dof_start + i]
-
-            # w = joint_qd[dof_start + i]
-            # r = joint_q[coord_start + i]
-
-            tau[dof_start + i] = -wp.dot(S_s, body_f_s)  # - w * target_kd - r * target_ke
-
-        return
-
-    if type == wp.sim.JOINT_FREE or type == wp.sim.JOINT_DISTANCE:
-        for i in range(6):
-            S_s = joint_S_s[dof_start + i]
-            tau[dof_start + i] = -wp.dot(S_s, body_f_s)
-
-        return
-
-    if type == wp.sim.JOINT_COMPOUND or type == wp.sim.JOINT_UNIVERSAL or type == wp.sim.JOINT_D6:
-        axis_count = lin_axis_count + ang_axis_count
-
-        for i in range(axis_count):
-            S_s = joint_S_s[dof_start + i]
-
-            q = joint_q[coord_start + i]
-            qd = joint_qd[dof_start + i]
-            act = joint_act[axis_start + i]
-
-            lower = joint_limit_lower[axis_start + i]
-            upper = joint_limit_upper[axis_start + i]
-            limit_ke = joint_limit_ke[axis_start + i]
-            limit_kd = joint_limit_kd[axis_start + i]
-            target_ke = joint_target_ke[axis_start + i]
-            target_kd = joint_target_kd[axis_start + i]
-            mode = joint_axis_mode[axis_start + i]
-
-            f = eval_joint_force(q, qd, act, target_ke, target_kd, lower, upper, limit_ke, limit_kd, mode)
-
-            # total torque / force on the joint
-            t = -wp.dot(S_s, body_f_s) + f
-
-            tau[dof_start + i] = t
-
-        return
-
-
-@wp.func
-def jcalc_integrate(
-    type: int,
-    joint_q: wp.array(dtype=float),
-    joint_qd: wp.array(dtype=float),
-    joint_qdd: wp.array(dtype=float),
-    coord_start: int,
-    dof_start: int,
-    lin_axis_count: int,
-    ang_axis_count: int,
-    dt: float,
-    # outputs
-    joint_q_new: wp.array(dtype=float),
-    joint_qd_new: wp.array(dtype=float),
-):
-    if type == wp.sim.JOINT_FIXED:
-        return
-
-    # prismatic / revolute
-    if type == wp.sim.JOINT_PRISMATIC or type == wp.sim.JOINT_REVOLUTE:
-        qdd = joint_qdd[dof_start]
-        qd = joint_qd[dof_start]
-        q = joint_q[coord_start]
-
-        qd_new = qd + qdd * dt
-        q_new = q + qd_new * dt
-
-        joint_qd_new[dof_start] = qd_new
-        joint_q_new[coord_start] = q_new
-
-        return
-
-    # ball
-    if type == wp.sim.JOINT_BALL:
-        m_j = wp.vec3(joint_qdd[dof_start + 0], joint_qdd[dof_start + 1], joint_qdd[dof_start + 2])
-        w_j = wp.vec3(joint_qd[dof_start + 0], joint_qd[dof_start + 1], joint_qd[dof_start + 2])
-
-        r_j = wp.quat(
-            joint_q[coord_start + 0], joint_q[coord_start + 1], joint_q[coord_start + 2], joint_q[coord_start + 3]
-        )
-
-        # symplectic Euler
-        w_j_new = w_j + m_j * dt
-
-        drdt_j = wp.quat(w_j_new, 0.0) * r_j * 0.5
-
-        # new orientation (normalized)
-        r_j_new = wp.normalize(r_j + drdt_j * dt)
-
-        # update joint coords
-        joint_q_new[coord_start + 0] = r_j_new[0]
-        joint_q_new[coord_start + 1] = r_j_new[1]
-        joint_q_new[coord_start + 2] = r_j_new[2]
-        joint_q_new[coord_start + 3] = r_j_new[3]
-
-        # update joint vel
-        joint_qd_new[dof_start + 0] = w_j_new[0]
-        joint_qd_new[dof_start + 1] = w_j_new[1]
-        joint_qd_new[dof_start + 2] = w_j_new[2]
-
-        return
-
-    # free joint
-    if type == wp.sim.JOINT_FREE or type == wp.sim.JOINT_DISTANCE:
-        # dofs: qd = (omega_x, omega_y, omega_z, vel_x, vel_y, vel_z)
-        # coords: q = (trans_x, trans_y, trans_z, quat_x, quat_y, quat_z, quat_w)
-
-        # angular and linear acceleration
-        m_s = wp.vec3(joint_qdd[dof_start + 0], joint_qdd[dof_start + 1], joint_qdd[dof_start + 2])
-        a_s = wp.vec3(joint_qdd[dof_start + 3], joint_qdd[dof_start + 4], joint_qdd[dof_start + 5])
-
-        # angular and linear velocity
-        w_s = wp.vec3(joint_qd[dof_start + 0], joint_qd[dof_start + 1], joint_qd[dof_start + 2])
-        v_s = wp.vec3(joint_qd[dof_start + 3], joint_qd[dof_start + 4], joint_qd[dof_start + 5])
-
-        # symplectic Euler
-        w_s = w_s + m_s * dt
-        v_s = v_s + a_s * dt
-
-        # translation of origin
-        p_s = wp.vec3(joint_q[coord_start + 0], joint_q[coord_start + 1], joint_q[coord_start + 2])
-
-        # linear vel of origin (note q/qd switch order of linear angular elements)
-        # note we are converting the body twist in the space frame (w_s, v_s) to compute center of mass velocity
-        dpdt_s = v_s + wp.cross(w_s, p_s)
-
-        # quat and quat derivative
-        r_s = wp.quat(
-            joint_q[coord_start + 3], joint_q[coord_start + 4], joint_q[coord_start + 5], joint_q[coord_start + 6]
-        )
-
-        drdt_s = wp.quat(w_s, 0.0) * r_s * 0.5
-
-        # new orientation (normalized)
-        p_s_new = p_s + dpdt_s * dt
-        r_s_new = wp.normalize(r_s + drdt_s * dt)
-
-        # update transform
-        joint_q_new[coord_start + 0] = p_s_new[0]
-        joint_q_new[coord_start + 1] = p_s_new[1]
-        joint_q_new[coord_start + 2] = p_s_new[2]
-
-        joint_q_new[coord_start + 3] = r_s_new[0]
-        joint_q_new[coord_start + 4] = r_s_new[1]
-        joint_q_new[coord_start + 5] = r_s_new[2]
-        joint_q_new[coord_start + 6] = r_s_new[3]
-
-        # update joint_twist
-        joint_qd_new[dof_start + 0] = w_s[0]
-        joint_qd_new[dof_start + 1] = w_s[1]
-        joint_qd_new[dof_start + 2] = w_s[2]
-        joint_qd_new[dof_start + 3] = v_s[0]
-        joint_qd_new[dof_start + 4] = v_s[1]
-        joint_qd_new[dof_start + 5] = v_s[2]
-
-        return
-
-    # other joint types (compound, universal, D6)
-    if type == wp.sim.JOINT_COMPOUND or type == wp.sim.JOINT_UNIVERSAL or type == wp.sim.JOINT_D6:
-        axis_count = lin_axis_count + ang_axis_count
-
-        for i in range(axis_count):
-            qdd = joint_qdd[dof_start + i]
-            qd = joint_qd[dof_start + i]
-            q = joint_q[coord_start + i]
-
-            qd_new = qd + qdd * dt
-            q_new = q + qd_new * dt
-
-            joint_qd_new[dof_start + i] = qd_new
-            joint_q_new[coord_start + i] = q_new
-
-        return
-
-
-@wp.func
-def compute_link_transform(
-    i: int,
-    joint_type: wp.array(dtype=int),
-    joint_parent: wp.array(dtype=int),
-    joint_child: wp.array(dtype=int),
-    joint_q_start: wp.array(dtype=int),
-    joint_q: wp.array(dtype=float),
-    joint_X_p: wp.array(dtype=wp.transform),
-    joint_X_c: wp.array(dtype=wp.transform),
-    body_X_com: 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),
-    # outputs
-    body_q: wp.array(dtype=wp.transform),
-    body_q_com: wp.array(dtype=wp.transform),
-):
-    # parent transform
-    parent = joint_parent[i]
-    child = joint_child[i]
-
-    # parent transform in spatial coordinates
-    X_pj = joint_X_p[i]
-    X_cj = joint_X_c[i]
-    # parent anchor frame in world space
-    X_wpj = X_pj
-    if parent >= 0:
-        X_wp = body_q[parent]
-        X_wpj = X_wp * X_wpj
-
-    type = joint_type[i]
-    axis_start = joint_axis_start[i]
-    lin_axis_count = joint_axis_dim[i, 0]
-    ang_axis_count = joint_axis_dim[i, 1]
-    coord_start = joint_q_start[i]
-
-    # compute transform across joint
-    X_j = jcalc_transform(type, joint_axis, axis_start, lin_axis_count, ang_axis_count, joint_q, coord_start)
-
-    # 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)
-
-    # compute transform of center of mass
-    X_cm = body_X_com[child]
-    X_sm = X_wc * X_cm
-
-    # store geometry transforms
-    body_q[child] = X_wc
-    body_q_com[child] = X_sm
-
-
-@wp.kernel
-def eval_rigid_fk(
-    articulation_start: wp.array(dtype=int),
-    joint_type: wp.array(dtype=int),
-    joint_parent: wp.array(dtype=int),
-    joint_child: wp.array(dtype=int),
-    joint_q_start: wp.array(dtype=int),
-    joint_q: wp.array(dtype=float),
-    joint_X_p: wp.array(dtype=wp.transform),
-    joint_X_c: wp.array(dtype=wp.transform),
-    body_X_com: 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),
-    # outputs
-    body_q: wp.array(dtype=wp.transform),
-    body_q_com: wp.array(dtype=wp.transform),
-):
-    # one thread per-articulation
-    index = wp.tid()
-
-    start = articulation_start[index]
-    end = articulation_start[index + 1]
-
-    for i in range(start, end):
-        compute_link_transform(
-            i,
-            joint_type,
-            joint_parent,
-            joint_child,
-            joint_q_start,
-            joint_q,
-            joint_X_p,
-            joint_X_c,
-            body_X_com,
-            joint_axis,
-            joint_axis_start,
-            joint_axis_dim,
-            body_q,
-            body_q_com,
-        )
-
-
-@wp.func
-def spatial_cross(a: wp.spatial_vector, b: wp.spatial_vector):
-    w_a = wp.spatial_top(a)
-    v_a = wp.spatial_bottom(a)
-
-    w_b = wp.spatial_top(b)
-    v_b = wp.spatial_bottom(b)
-
-    w = wp.cross(w_a, w_b)
-    v = wp.cross(w_a, v_b) + wp.cross(v_a, w_b)
-
-    return wp.spatial_vector(w, v)
-
-
-@wp.func
-def spatial_cross_dual(a: wp.spatial_vector, b: wp.spatial_vector):
-    w_a = wp.spatial_top(a)
-    v_a = wp.spatial_bottom(a)
-
-    w_b = wp.spatial_top(b)
-    v_b = wp.spatial_bottom(b)
-
-    w = wp.cross(w_a, w_b) + wp.cross(v_a, v_b)
-    v = wp.cross(w_a, v_b)
-
-    return wp.spatial_vector(w, v)
-
-
-@wp.func
-def dense_index(stride: int, i: int, j: int):
-    return i * stride + j
-
-
-@wp.func
-def compute_link_velocity(
-    i: int,
-    joint_type: wp.array(dtype=int),
-    joint_parent: wp.array(dtype=int),
-    joint_child: wp.array(dtype=int),
-    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),
-    joint_axis: wp.array(dtype=wp.vec3),
-    joint_axis_start: wp.array(dtype=int),
-    joint_axis_dim: wp.array(dtype=int, ndim=2),
-    body_I_m: wp.array(dtype=wp.spatial_matrix),
-    body_q: wp.array(dtype=wp.transform),
-    body_q_com: wp.array(dtype=wp.transform),
-    joint_X_p: wp.array(dtype=wp.transform),
-    joint_X_c: wp.array(dtype=wp.transform),
-    gravity: wp.vec3,
-    # outputs
-    joint_S_s: wp.array(dtype=wp.spatial_vector),
-    body_I_s: wp.array(dtype=wp.spatial_matrix),
-    body_v_s: wp.array(dtype=wp.spatial_vector),
-    body_f_s: wp.array(dtype=wp.spatial_vector),
-    body_a_s: wp.array(dtype=wp.spatial_vector),
-):
-    type = joint_type[i]
-    child = joint_child[i]
-    parent = joint_parent[i]
-    q_start = joint_q_start[i]
-    qd_start = joint_qd_start[i]
-
-    X_pj = joint_X_p[i]
-    # X_cj = joint_X_c[i]
-
-    # parent anchor frame in world space
-    X_wpj = X_pj
-    if parent >= 0:
-        X_wp = body_q[parent]
-        X_wpj = X_wp * X_wpj
-
-    # compute motion subspace and velocity across the joint (also stores S_s to global memory)
-    axis_start = joint_axis_start[i]
-    lin_axis_count = joint_axis_dim[i, 0]
-    ang_axis_count = joint_axis_dim[i, 1]
-    v_j_s = jcalc_motion(
-        type,
-        joint_axis,
-        axis_start,
-        lin_axis_count,
-        ang_axis_count,
-        X_wpj,
-        joint_q,
-        joint_qd,
-        q_start,
-        qd_start,
-        joint_S_s,
-    )
-
-    # parent velocity
-    v_parent_s = wp.spatial_vector()
-    a_parent_s = wp.spatial_vector()
-
-    if parent >= 0:
-        v_parent_s = body_v_s[parent]
-        a_parent_s = body_a_s[parent]
-
-    # body velocity, acceleration
-    v_s = v_parent_s + v_j_s
-    a_s = a_parent_s + spatial_cross(v_s, v_j_s)  # + joint_S_s[i]*self.joint_qdd[i]
-
-    # compute body forces
-    X_sm = body_q_com[child]
-    I_m = body_I_m[child]
-
-    # gravity and external forces (expressed in frame aligned with s but centered at body mass)
-    m = I_m[3, 3]
-
-    f_g = m * gravity
-    r_com = wp.transform_get_translation(X_sm)
-    f_g_s = wp.spatial_vector(wp.cross(r_com, f_g), f_g)
-
-    # body forces
-    I_s = spatial_transform_inertia(X_sm, I_m)
-
-    f_b_s = I_s * a_s + spatial_cross_dual(v_s, I_s * v_s)
-
-    body_v_s[child] = v_s
-    body_a_s[child] = a_s
-    body_f_s[child] = f_b_s - f_g_s
-    body_I_s[child] = I_s
-
-
-# Inverse dynamics via Recursive Newton-Euler algorithm (Featherstone Table 5.1)
-@wp.kernel
-def eval_rigid_id(
-    articulation_start: wp.array(dtype=int),
-    joint_type: wp.array(dtype=int),
-    joint_parent: wp.array(dtype=int),
-    joint_child: wp.array(dtype=int),
-    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),
-    joint_axis: wp.array(dtype=wp.vec3),
-    joint_axis_start: wp.array(dtype=int),
-    joint_axis_dim: wp.array(dtype=int, ndim=2),
-    body_I_m: wp.array(dtype=wp.spatial_matrix),
-    body_q: wp.array(dtype=wp.transform),
-    body_q_com: wp.array(dtype=wp.transform),
-    joint_X_p: wp.array(dtype=wp.transform),
-    joint_X_c: wp.array(dtype=wp.transform),
-    gravity: wp.vec3,
-    # outputs
-    joint_S_s: wp.array(dtype=wp.spatial_vector),
-    body_I_s: wp.array(dtype=wp.spatial_matrix),
-    body_v_s: wp.array(dtype=wp.spatial_vector),
-    body_f_s: wp.array(dtype=wp.spatial_vector),
-    body_a_s: wp.array(dtype=wp.spatial_vector),
-):
-    # one thread per-articulation
-    index = wp.tid()
-
-    start = articulation_start[index]
-    end = articulation_start[index + 1]
-
-    # compute link velocities and coriolis forces
-    for i in range(start, end):
-        compute_link_velocity(
-            i,
-            joint_type,
-            joint_parent,
-            joint_child,
-            joint_q_start,
-            joint_qd_start,
-            joint_q,
-            joint_qd,
-            joint_axis,
-            joint_axis_start,
-            joint_axis_dim,
-            body_I_m,
-            body_q,
-            body_q_com,
-            joint_X_p,
-            joint_X_c,
-            gravity,
-            joint_S_s,
-            body_I_s,
-            body_v_s,
-            body_f_s,
-            body_a_s,
-        )
-
-
-@wp.kernel
-def eval_rigid_tau(
-    articulation_start: wp.array(dtype=int),
-    joint_type: wp.array(dtype=int),
-    joint_parent: wp.array(dtype=int),
-    joint_child: wp.array(dtype=int),
-    joint_q_start: wp.array(dtype=int),
-    joint_qd_start: wp.array(dtype=int),
-    joint_axis_start: wp.array(dtype=int),
-    joint_axis_dim: wp.array(dtype=int, ndim=2),
-    joint_axis_mode: wp.array(dtype=int),
-    joint_q: wp.array(dtype=float),
-    joint_qd: wp.array(dtype=float),
-    joint_act: wp.array(dtype=float),
-    joint_target_ke: wp.array(dtype=float),
-    joint_target_kd: wp.array(dtype=float),
-    joint_limit_lower: wp.array(dtype=float),
-    joint_limit_upper: wp.array(dtype=float),
-    joint_limit_ke: wp.array(dtype=float),
-    joint_limit_kd: wp.array(dtype=float),
-    joint_S_s: wp.array(dtype=wp.spatial_vector),
-    body_fb_s: wp.array(dtype=wp.spatial_vector),
-    body_f_ext: wp.array(dtype=wp.spatial_vector),
-    # outputs
-    body_ft_s: wp.array(dtype=wp.spatial_vector),
-    tau: wp.array(dtype=float),
-):
-    # one thread per-articulation
-    index = wp.tid()
-
-    start = articulation_start[index]
-    end = articulation_start[index + 1]
-    count = end - start
-
-    # compute joint forces
-    for offset in range(count):
-        # for backwards traversal
-        i = end - offset - 1
-
-        type = joint_type[i]
-        parent = joint_parent[i]
-        child = joint_child[i]
-        dof_start = joint_qd_start[i]
-        coord_start = joint_q_start[i]
-        axis_start = joint_axis_start[i]
-        lin_axis_count = joint_axis_dim[i, 0]
-        ang_axis_count = joint_axis_dim[i, 1]
-
-        # total forces on body
-        f_b_s = body_fb_s[child]
-        f_t_s = body_ft_s[child]
-        f_ext = body_f_ext[child]
-        f_s = f_b_s + f_t_s + f_ext
-
-        # compute joint-space forces, writes out tau
-        jcalc_tau(
-            type,
-            joint_target_ke,
-            joint_target_kd,
-            joint_limit_ke,
-            joint_limit_kd,
-            joint_S_s,
-            joint_q,
-            joint_qd,
-            joint_act,
-            joint_axis_mode,
-            joint_limit_lower,
-            joint_limit_upper,
-            coord_start,
-            dof_start,
-            axis_start,
-            lin_axis_count,
-            ang_axis_count,
-            f_s,
-            tau,
-        )
-
-        # update parent forces, todo: check that this is valid for the backwards pass
-        if parent >= 0:
-            wp.atomic_add(body_ft_s, parent, f_s)
-
-
-# builds spatial Jacobian J which is an (joint_count*6)x(dof_count) matrix
-@wp.kernel
-def eval_rigid_jacobian(
-    articulation_start: wp.array(dtype=int),
-    articulation_J_start: wp.array(dtype=int),
-    joint_ancestor: wp.array(dtype=int),
-    joint_qd_start: wp.array(dtype=int),
-    joint_S_s: wp.array(dtype=wp.spatial_vector),
-    # outputs
-    J: wp.array(dtype=float),
-):
-    # one thread per-articulation
-    index = wp.tid()
-
-    joint_start = articulation_start[index]
-    joint_end = articulation_start[index + 1]
-    joint_count = joint_end - joint_start
-
-    J_offset = articulation_J_start[index]
-
-    articulation_dof_start = joint_qd_start[joint_start]
-    articulation_dof_end = joint_qd_start[joint_end]
-    articulation_dof_count = articulation_dof_end - articulation_dof_start
-
-    for i in range(joint_count):
-        row_start = i * 6
-
-        j = joint_start + i
-        while j != -1:
-            joint_dof_start = joint_qd_start[j]
-            joint_dof_end = joint_qd_start[j + 1]
-            joint_dof_count = joint_dof_end - joint_dof_start
-
-            # fill out each row of the Jacobian walking up the tree
-            for dof in range(joint_dof_count):
-                col = (joint_dof_start - articulation_dof_start) + dof
-                S = joint_S_s[joint_dof_start + dof]
-
-                for k in range(6):
-                    J[J_offset + dense_index(articulation_dof_count, row_start + k, col)] = S[k]
-
-            j = joint_ancestor[j]
-
-
-@wp.func
-def spatial_mass(
-    body_I_s: wp.array(dtype=wp.spatial_matrix),
-    joint_start: int,
-    joint_count: int,
-    M_start: int,
-    # outputs
-    M: wp.array(dtype=float),
-):
-    stride = joint_count * 6
-    for l in range(joint_count):
-        I = body_I_s[joint_start + l]
-        for i in range(6):
-            for j in range(6):
-                M[M_start + dense_index(stride, l * 6 + i, l * 6 + j)] = I[i, j]
-
-
-@wp.kernel
-def eval_rigid_mass(
-    articulation_start: wp.array(dtype=int),
-    articulation_M_start: wp.array(dtype=int),
-    body_I_s: wp.array(dtype=wp.spatial_matrix),
-    # outputs
-    M: wp.array(dtype=float),
-):
-    # one thread per-articulation
-    index = wp.tid()
-
-    joint_start = articulation_start[index]
-    joint_end = articulation_start[index + 1]
-    joint_count = joint_end - joint_start
-
-    M_offset = articulation_M_start[index]
-
-    spatial_mass(body_I_s, joint_start, joint_count, M_offset, M)
-
-
-@wp.func
-def dense_gemm(
-    m: int,
-    n: int,
-    p: int,
-    transpose_A: bool,
-    transpose_B: bool,
-    add_to_C: bool,
-    A_start: int,
-    B_start: int,
-    C_start: int,
-    A: wp.array(dtype=float),
-    B: wp.array(dtype=float),
-    # outputs
-    C: wp.array(dtype=float),
-):
-    # multiply a `m x p` matrix A by a `p x n` matrix B to produce a `m x n` matrix C
-    for i in range(m):
-        for j in range(n):
-            sum = float(0.0)
-            for k in range(p):
-                if transpose_A:
-                    a_i = k * m + i
-                else:
-                    a_i = i * p + k
-                if transpose_B:
-                    b_j = j * p + k
-                else:
-                    b_j = k * n + j
-                sum += A[A_start + a_i] * B[B_start + b_j]
-
-            if add_to_C:
-                C[C_start + i * n + j] += sum
-            else:
-                C[C_start + i * n + j] = sum
-
-
-# @wp.func_grad(dense_gemm)
-# def adj_dense_gemm(
-#     m: int,
-#     n: int,
-#     p: int,
-#     transpose_A: bool,
-#     transpose_B: bool,
-#     add_to_C: bool,
-#     A_start: int,
-#     B_start: int,
-#     C_start: int,
-#     A: wp.array(dtype=float),
-#     B: wp.array(dtype=float),
-#     # outputs
-#     C: wp.array(dtype=float),
-# ):
-#     add_to_C = True
-#     if transpose_A:
-#         dense_gemm(p, m, n, False, True, add_to_C, A_start, B_start, C_start, B, wp.adjoint[C], wp.adjoint[A])
-#         dense_gemm(p, n, m, False, False, add_to_C, A_start, B_start, C_start, A, wp.adjoint[C], wp.adjoint[B])
-#     else:
-#         dense_gemm(
-#             m, p, n, False, not transpose_B, add_to_C, A_start, B_start, C_start, wp.adjoint[C], B, wp.adjoint[A]
-#         )
-#         dense_gemm(p, n, m, True, False, add_to_C, A_start, B_start, C_start, A, wp.adjoint[C], wp.adjoint[B])
-
-
-def create_inertia_matrix_kernel(num_joints, num_dofs):
-    @wp.kernel
-    def eval_dense_gemm_tile(
-        J_arr: wp.array3d(dtype=float), M_arr: wp.array3d(dtype=float), H_arr: wp.array3d(dtype=float)
-    ):
-        articulation = wp.tid()
-
-        J = wp.tile_load(J_arr[articulation], shape=(wp.static(6 * num_joints), num_dofs))
-        P = wp.tile_zeros(shape=(wp.static(6 * num_joints), num_dofs), dtype=float)
-
-        # compute P = M*J where M is a 6x6 block diagonal mass matrix
-        for i in range(int(num_joints)):
-            # 6x6 block matrices are on the diagonal
-            M_body = wp.tile_load(M_arr[articulation], shape=(6, 6), offset=(i * 6, i * 6))
-
-            # load a 6xN row from the Jacobian
-            J_body = wp.tile_view(J, offset=(i * 6, 0), shape=(6, num_dofs))
-
-            # compute weighted row
-            P_body = wp.tile_matmul(M_body, J_body)
-
-            # assign to the P slice
-            wp.tile_assign(P, P_body, offset=(i * 6, 0))
-
-        # compute H = J^T*P
-        H = wp.tile_matmul(wp.tile_transpose(J), P)
-
-        wp.tile_store(H_arr[articulation], H)
-
-    return eval_dense_gemm_tile
-
-
-def create_batched_cholesky_kernel(num_dofs):
-    assert num_dofs == 18
-
-    @wp.kernel
-    def eval_tiled_dense_cholesky_batched(
-        A: wp.array3d(dtype=float), R: wp.array2d(dtype=float), L: wp.array3d(dtype=float)
-    ):
-        articulation = wp.tid()
-
-        a = wp.tile_load(A[articulation], shape=(num_dofs, num_dofs), storage="shared")
-        r = wp.tile_load(R[articulation], shape=num_dofs, storage="shared")
-        a_r = wp.tile_diag_add(a, r)
-        l = wp.tile_cholesky(a_r)
-        wp.tile_store(L[articulation], wp.tile_transpose(l))
-
-    return eval_tiled_dense_cholesky_batched
-
-
-def create_inertia_matrix_cholesky_kernel(num_joints, num_dofs):
-    @wp.kernel
-    def eval_dense_gemm_and_cholesky_tile(
-        J_arr: wp.array3d(dtype=float),
-        M_arr: wp.array3d(dtype=float),
-        R_arr: wp.array2d(dtype=float),
-        H_arr: wp.array3d(dtype=float),
-        L_arr: wp.array3d(dtype=float),
-    ):
-        articulation = wp.tid()
-
-        J = wp.tile_load(J_arr[articulation], shape=(wp.static(6 * num_joints), num_dofs))
-        P = wp.tile_zeros(shape=(wp.static(6 * num_joints), num_dofs), dtype=float)
-
-        # compute P = M*J where M is a 6x6 block diagonal mass matrix
-        for i in range(int(num_joints)):
-            # 6x6 block matrices are on the diagonal
-            M_body = wp.tile_load(M_arr[articulation], shape=(6, 6), offset=(i * 6, i * 6))
-
-            # load a 6xN row from the Jacobian
-            J_body = wp.tile_view(J, offset=(i * 6, 0), shape=(6, num_dofs))
-
-            # compute weighted row
-            P_body = wp.tile_matmul(M_body, J_body)
-
-            # assign to the P slice
-            wp.tile_assign(P, P_body, offset=(i * 6, 0))
-
-        # compute H = J^T*P
-        H = wp.tile_matmul(wp.tile_transpose(J), P)
-        wp.tile_store(H_arr[articulation], H)
-
-        # cholesky L L^T = (H + diag(R))
-        R = wp.tile_load(R_arr[articulation], shape=num_dofs, storage="shared")
-        H_R = wp.tile_diag_add(H, R)
-        L = wp.tile_cholesky(H_R)
-        wp.tile_store(L_arr[articulation], L)
-
-    return eval_dense_gemm_and_cholesky_tile
-
-
-@wp.kernel
-def eval_dense_gemm_batched(
-    m: wp.array(dtype=int),
-    n: wp.array(dtype=int),
-    p: wp.array(dtype=int),
-    transpose_A: bool,
-    transpose_B: bool,
-    A_start: wp.array(dtype=int),
-    B_start: wp.array(dtype=int),
-    C_start: wp.array(dtype=int),
-    A: wp.array(dtype=float),
-    B: wp.array(dtype=float),
-    C: wp.array(dtype=float),
-):
-    # on the CPU each thread computes the whole matrix multiply
-    # on the GPU each block computes the multiply with one output per-thread
-    batch = wp.tid()  # /kNumThreadsPerBlock;
-    add_to_C = False
-
-    dense_gemm(
-        m[batch],
-        n[batch],
-        p[batch],
-        transpose_A,
-        transpose_B,
-        add_to_C,
-        A_start[batch],
-        B_start[batch],
-        C_start[batch],
-        A,
-        B,
-        C,
-    )
-
-
-@wp.func
-def dense_cholesky(
-    n: int,
-    A: wp.array(dtype=float),
-    R: wp.array(dtype=float),
-    A_start: int,
-    R_start: int,
-    # outputs
-    L: wp.array(dtype=float),
-):
-    # compute the Cholesky factorization of A = L L^T with diagonal regularization R
-    for j in range(n):
-        s = A[A_start + dense_index(n, j, j)] + R[R_start + j]
-
-        for k in range(j):
-            r = L[A_start + dense_index(n, j, k)]
-            s -= r * r
-
-        s = wp.sqrt(s)
-        invS = 1.0 / s
-
-        L[A_start + dense_index(n, j, j)] = s
-
-        for i in range(j + 1, n):
-            s = A[A_start + dense_index(n, i, j)]
-
-            for k in range(j):
-                s -= L[A_start + dense_index(n, i, k)] * L[A_start + dense_index(n, j, k)]
-
-            L[A_start + dense_index(n, i, j)] = s * invS
-
-
-@wp.func_grad(dense_cholesky)
-def adj_dense_cholesky(
-    n: int,
-    A: wp.array(dtype=float),
-    R: wp.array(dtype=float),
-    A_start: int,
-    R_start: int,
-    # outputs
-    L: wp.array(dtype=float),
-):
-    # nop, use dense_solve to differentiate through (A^-1)b = x
-    pass
-
-
-@wp.kernel
-def eval_dense_cholesky_batched(
-    A_starts: wp.array(dtype=int),
-    A_dim: wp.array(dtype=int),
-    A: wp.array(dtype=float),
-    R: wp.array(dtype=float),
-    L: wp.array(dtype=float),
-):
-    batch = wp.tid()
-
-    n = A_dim[batch]
-    A_start = A_starts[batch]
-    R_start = n * batch
-
-    dense_cholesky(n, A, R, A_start, R_start, L)
-
-
-@wp.func
-def dense_subs(
-    n: int,
-    L_start: int,
-    b_start: int,
-    L: wp.array(dtype=float),
-    b: wp.array(dtype=float),
-    # outputs
-    x: wp.array(dtype=float),
-):
-    # Solves (L L^T) x = b for x given the Cholesky factor L
-    # forward substitution solves the lower triangular system L y = b for y
-    for i in range(n):
-        s = b[b_start + i]
-
-        for j in range(i):
-            s -= L[L_start + dense_index(n, i, j)] * x[b_start + j]
-
-        x[b_start + i] = s / L[L_start + dense_index(n, i, i)]
-
-    # backward substitution solves the upper triangular system L^T x = y for x
-    for i in range(n - 1, -1, -1):
-        s = x[b_start + i]
-
-        for j in range(i + 1, n):
-            s -= L[L_start + dense_index(n, j, i)] * x[b_start + j]
-
-        x[b_start + i] = s / L[L_start + dense_index(n, i, i)]
-
-
-@wp.func
-def dense_solve(
-    n: int,
-    L_start: int,
-    b_start: int,
-    A: wp.array(dtype=float),
-    L: wp.array(dtype=float),
-    b: wp.array(dtype=float),
-    # outputs
-    x: wp.array(dtype=float),
-    tmp: wp.array(dtype=float),
-):
-    # helper function to include tmp argument for backward pass
-    dense_subs(n, L_start, b_start, L, b, x)
-
-
-@wp.func_grad(dense_solve)
-def adj_dense_solve(
-    n: int,
-    L_start: int,
-    b_start: int,
-    A: wp.array(dtype=float),
-    L: wp.array(dtype=float),
-    b: wp.array(dtype=float),
-    # outputs
-    x: wp.array(dtype=float),
-    tmp: wp.array(dtype=float),
-):
-    if not tmp or not wp.adjoint[x] or not wp.adjoint[A] or not wp.adjoint[L]:
-        return
-    for i in range(n):
-        tmp[b_start + i] = 0.0
-
-    dense_subs(n, L_start, b_start, L, wp.adjoint[x], tmp)
-
-    for i in range(n):
-        wp.adjoint[b][b_start + i] += tmp[b_start + i]
-
-    # A* = -adj_b*x^T
-    for i in range(n):
-        for j in range(n):
-            wp.adjoint[L][L_start + dense_index(n, i, j)] += -tmp[b_start + i] * x[b_start + j]
-
-    for i in range(n):
-        for j in range(n):
-            wp.adjoint[A][L_start + dense_index(n, i, j)] += -tmp[b_start + i] * x[b_start + j]
-
-
-@wp.kernel
-def eval_dense_solve_batched(
-    L_start: wp.array(dtype=int),
-    L_dim: wp.array(dtype=int),
-    b_start: wp.array(dtype=int),
-    A: wp.array(dtype=float),
-    L: wp.array(dtype=float),
-    b: wp.array(dtype=float),
-    # outputs
-    x: wp.array(dtype=float),
-    tmp: wp.array(dtype=float),
-):
-    batch = wp.tid()
-
-    dense_solve(L_dim[batch], L_start[batch], b_start[batch], A, L, b, x, tmp)
-
-
-@wp.kernel
-def integrate_generalized_joints(
-    joint_type: wp.array(dtype=int),
-    joint_q_start: wp.array(dtype=int),
-    joint_qd_start: wp.array(dtype=int),
-    joint_axis_dim: wp.array(dtype=int, ndim=2),
-    joint_q: wp.array(dtype=float),
-    joint_qd: wp.array(dtype=float),
-    joint_qdd: wp.array(dtype=float),
-    dt: float,
-    # outputs
-    joint_q_new: wp.array(dtype=float),
-    joint_qd_new: wp.array(dtype=float),
-):
-    # one thread per-articulation
-    index = wp.tid()
-
-    type = joint_type[index]
-    coord_start = joint_q_start[index]
-    dof_start = joint_qd_start[index]
-    lin_axis_count = joint_axis_dim[index, 0]
-    ang_axis_count = joint_axis_dim[index, 1]
-
-    jcalc_integrate(
-        type,
-        joint_q,
-        joint_qd,
-        joint_qdd,
-        coord_start,
-        dof_start,
-        lin_axis_count,
-        ang_axis_count,
-        dt,
-        joint_q_new,
-        joint_qd_new,
-    )
-
-
-class FeatherstoneIntegrator(Integrator):
-    """A semi-implicit integrator using symplectic Euler that operates
-    on reduced (also called generalized) coordinates to simulate articulated rigid body dynamics
-    based on Featherstone's composite rigid body algorithm (CRBA).
-
-    See: Featherstone, Roy. Rigid Body Dynamics Algorithms. Springer US, 2014.
-
-    Instead of maximal coordinates :attr:`State.body_q` (rigid body positions) and :attr:`State.body_qd`
-    (rigid body velocities) as is the case :class:`SemiImplicitIntegrator`, :class:`FeatherstoneIntegrator`
-    uses :attr:`State.joint_q` and :attr:`State.joint_qd` to represent the positions and velocities of
-    joints without allowing any redundant degrees of freedom.
-
-    After constructing :class:`Model` and :class:`State` objects this time-integrator
-    may be used to advance the simulation state forward in time.
-
-    Note:
-        Unlike :class:`SemiImplicitIntegrator` and :class:`XPBDIntegrator`, :class:`FeatherstoneIntegrator` does not simulate rigid bodies with nonzero mass as floating bodies if they are not connected through any joints. Floating-base systems require an explicit free joint with which the body is connected to the world, see :meth:`ModelBuilder.add_joint_free`.
-
-    Semi-implicit time integration is a variational integrator that
-    preserves energy, however it not unconditionally stable, and requires a time-step
-    small enough to support the required stiffness and damping forces.
-
-    See: https://en.wikipedia.org/wiki/Semi-implicit_Euler_method
-
-    Example
-    -------
-
-    .. code-block:: python
-
-        integrator = wp.FeatherstoneIntegrator(model)
-
-        # simulation loop
-        for i in range(100):
-            state = integrator.simulate(model, state_in, state_out, dt)
-
-    Note:
-        The :class:`FeatherstoneIntegrator` requires the :class:`Model` to be passed in as a constructor argument.
-
-    """
-
-    def __init__(
-        self,
-        model,
-        angular_damping=0.05,
-        update_mass_matrix_every=1,
-        friction_smoothing=1.0,
-        use_tile_gemm=False,
-        fuse_cholesky=True,
-    ):
-        """
-        Args:
-            model (Model): the model to be simulated.
-            angular_damping (float, optional): Angular damping factor. Defaults to 0.05.
-            update_mass_matrix_every (int, optional): How often to update the mass matrix (every n-th time the :meth:`simulate` function gets called). Defaults to 1.
-            friction_smoothing (float, optional): The delta value for the Huber norm (see :func:`warp.math.norm_huber`) used for the friction velocity normalization. Defaults to 1.0.
-        """
-        self.angular_damping = angular_damping
-        self.update_mass_matrix_every = update_mass_matrix_every
-        self.friction_smoothing = friction_smoothing
-        self.use_tile_gemm = use_tile_gemm
-        self.fuse_cholesky = fuse_cholesky
-
-        self._step = 0
-
-        self.compute_articulation_indices(model)
-        self.allocate_model_aux_vars(model)
-
-        if self.use_tile_gemm:
-            # create a custom kernel to evaluate the system matrix for this type
-            if self.fuse_cholesky:
-                self.eval_inertia_matrix_cholesky_kernel = create_inertia_matrix_cholesky_kernel(
-                    int(self.joint_count), int(self.dof_count)
-                )
-            else:
-                self.eval_inertia_matrix_kernel = create_inertia_matrix_kernel(
-                    int(self.joint_count), int(self.dof_count)
-                )
-
-            # ensure matrix is reloaded since otherwise an unload can happen during graph capture
-            # todo: should not be necessary?
-            wp.load_module(device=wp.get_device())
-
-    def compute_articulation_indices(self, model):
-        # calculate total size and offsets of Jacobian and mass matrices for entire system
-        if model.joint_count:
-            self.J_size = 0
-            self.M_size = 0
-            self.H_size = 0
-
-            articulation_J_start = []
-            articulation_M_start = []
-            articulation_H_start = []
-
-            articulation_M_rows = []
-            articulation_H_rows = []
-            articulation_J_rows = []
-            articulation_J_cols = []
-
-            articulation_dof_start = []
-            articulation_coord_start = []
-
-            articulation_start = model.articulation_start.numpy()
-            joint_q_start = model.joint_q_start.numpy()
-            joint_qd_start = model.joint_qd_start.numpy()
-
-            for i in range(model.articulation_count):
-                first_joint = articulation_start[i]
-                last_joint = articulation_start[i + 1]
-
-                first_coord = joint_q_start[first_joint]
-
-                first_dof = joint_qd_start[first_joint]
-                last_dof = joint_qd_start[last_joint]
-
-                joint_count = last_joint - first_joint
-                dof_count = last_dof - first_dof
-
-                articulation_J_start.append(self.J_size)
-                articulation_M_start.append(self.M_size)
-                articulation_H_start.append(self.H_size)
-                articulation_dof_start.append(first_dof)
-                articulation_coord_start.append(first_coord)
-
-                # bit of data duplication here, but will leave it as such for clarity
-                articulation_M_rows.append(joint_count * 6)
-                articulation_H_rows.append(dof_count)
-                articulation_J_rows.append(joint_count * 6)
-                articulation_J_cols.append(dof_count)
-
-                if self.use_tile_gemm:
-                    # store the joint and dof count assuming all
-                    # articulations have the same structure
-                    self.joint_count = joint_count
-                    self.dof_count = dof_count
-
-                self.J_size += 6 * joint_count * dof_count
-                self.M_size += 6 * joint_count * 6 * joint_count
-                self.H_size += dof_count * dof_count
-
-            # matrix offsets for batched gemm
-            self.articulation_J_start = wp.array(articulation_J_start, dtype=wp.int32, device=model.device)
-            self.articulation_M_start = wp.array(articulation_M_start, dtype=wp.int32, device=model.device)
-            self.articulation_H_start = wp.array(articulation_H_start, dtype=wp.int32, device=model.device)
-
-            self.articulation_M_rows = wp.array(articulation_M_rows, dtype=wp.int32, device=model.device)
-            self.articulation_H_rows = wp.array(articulation_H_rows, dtype=wp.int32, device=model.device)
-            self.articulation_J_rows = wp.array(articulation_J_rows, dtype=wp.int32, device=model.device)
-            self.articulation_J_cols = wp.array(articulation_J_cols, dtype=wp.int32, device=model.device)
-
-            self.articulation_dof_start = wp.array(articulation_dof_start, dtype=wp.int32, device=model.device)
-            self.articulation_coord_start = wp.array(articulation_coord_start, dtype=wp.int32, device=model.device)
-
-    def allocate_model_aux_vars(self, model):
-        # allocate mass, Jacobian matrices, and other auxiliary variables pertaining to the model
-        if model.joint_count:
-            # system matrices
-            self.M = wp.zeros((self.M_size,), dtype=wp.float32, device=model.device, requires_grad=model.requires_grad)
-            self.J = wp.zeros((self.J_size,), dtype=wp.float32, device=model.device, requires_grad=model.requires_grad)
-            self.P = wp.empty_like(self.J, requires_grad=model.requires_grad)
-            self.H = wp.empty((self.H_size,), dtype=wp.float32, device=model.device, requires_grad=model.requires_grad)
-
-            # zero since only upper triangle is set which can trigger NaN detection
-            self.L = wp.zeros_like(self.H)
-
-        if model.body_count:
-            self.body_I_m = wp.empty(
-                (model.body_count,), dtype=wp.spatial_matrix, device=model.device, requires_grad=model.requires_grad
-            )
-            wp.launch(
-                compute_spatial_inertia,
-                model.body_count,
-                inputs=[model.body_inertia, model.body_mass],
-                outputs=[self.body_I_m],
-                device=model.device,
-            )
-            self.body_X_com = wp.empty(
-                (model.body_count,), dtype=wp.transform, device=model.device, requires_grad=model.requires_grad
-            )
-            wp.launch(
-                compute_com_transforms,
-                model.body_count,
-                inputs=[model.body_com],
-                outputs=[self.body_X_com],
-                device=model.device,
-            )
-
-    def allocate_state_aux_vars(self, model, target, requires_grad):
-        # allocate auxiliary variables that vary with state
-        if model.body_count:
-            # joints
-            target.joint_qdd = wp.zeros_like(model.joint_qd, requires_grad=requires_grad)
-            target.joint_tau = wp.empty_like(model.joint_qd, requires_grad=requires_grad)
-            if requires_grad:
-                # used in the custom grad implementation of eval_dense_solve_batched
-                target.joint_solve_tmp = wp.zeros_like(model.joint_qd, requires_grad=True)
-            else:
-                target.joint_solve_tmp = None
-            target.joint_S_s = wp.empty(
-                (model.joint_dof_count,),
-                dtype=wp.spatial_vector,
-                device=model.device,
-                requires_grad=requires_grad,
-            )
-
-            # derived rigid body data (maximal coordinates)
-            target.body_q_com = wp.empty_like(model.body_q, requires_grad=requires_grad)
-            target.body_I_s = wp.empty(
-                (model.body_count,), dtype=wp.spatial_matrix, device=model.device, requires_grad=requires_grad
-            )
-            target.body_v_s = wp.empty(
-                (model.body_count,), dtype=wp.spatial_vector, device=model.device, requires_grad=requires_grad
-            )
-            target.body_a_s = wp.empty(
-                (model.body_count,), dtype=wp.spatial_vector, device=model.device, requires_grad=requires_grad
-            )
-            target.body_f_s = wp.zeros(
-                (model.body_count,), dtype=wp.spatial_vector, device=model.device, requires_grad=requires_grad
-            )
-            target.body_ft_s = wp.zeros(
-                (model.body_count,), dtype=wp.spatial_vector, device=model.device, requires_grad=requires_grad
-            )
-
-            target._featherstone_augmented = True
-
-    def simulate(self, model: Model, state_in: State, state_out: State, dt: float, control: Control = None):
-        requires_grad = state_in.requires_grad
-
-        # optionally create dynamical auxiliary variables
-        if requires_grad:
-            state_aug = state_out
-        else:
-            state_aug = self
-
-        if not getattr(state_aug, "_featherstone_augmented", False):
-            self.allocate_state_aux_vars(model, state_aug, requires_grad)
-        if control is None:
-            control = model.control(clone_variables=False)
-
-        with wp.ScopedTimer("simulate", False):
-            particle_f = None
-            body_f = None
-
-            if state_in.particle_count:
-                particle_f = state_in.particle_f
-
-            if state_in.body_count:
-                body_f = state_in.body_f
-
-            # damped springs
-            eval_spring_forces(model, state_in, particle_f)
-
-            # triangle elastic and lift/drag forces
-            eval_triangle_forces(model, state_in, control, particle_f)
-
-            # triangle/triangle contacts
-            eval_triangle_contact_forces(model, state_in, particle_f)
-
-            # triangle bending
-            eval_bending_forces(model, state_in, particle_f)
-
-            # tetrahedral FEM
-            eval_tetrahedral_forces(model, state_in, control, particle_f)
-
-            # particle-particle interactions
-            eval_particle_forces(model, state_in, particle_f)
-
-            # particle ground contacts
-            eval_particle_ground_contact_forces(model, state_in, particle_f)
-
-            # particle shape contact
-            eval_particle_body_contact_forces(model, state_in, particle_f, body_f, body_f_in_world_frame=True)
-
-            # muscles
-            if False:
-                eval_muscle_forces(model, state_in, control, body_f)
-
-            # ----------------------------
-            # articulations
-
-            if model.joint_count:
-                # evaluate body transforms
-                wp.launch(
-                    eval_rigid_fk,
-                    dim=model.articulation_count,
-                    inputs=[
-                        model.articulation_start,
-                        model.joint_type,
-                        model.joint_parent,
-                        model.joint_child,
-                        model.joint_q_start,
-                        state_in.joint_q,
-                        model.joint_X_p,
-                        model.joint_X_c,
-                        self.body_X_com,
-                        model.joint_axis,
-                        model.joint_axis_start,
-                        model.joint_axis_dim,
-                    ],
-                    outputs=[state_in.body_q, state_aug.body_q_com],
-                    device=model.device,
-                )
-
-                # print("body_X_sc:")
-                # print(state_in.body_q.numpy())
-
-                # evaluate joint inertias, motion vectors, and forces
-                state_aug.body_f_s.zero_()
-                wp.launch(
-                    eval_rigid_id,
-                    dim=model.articulation_count,
-                    inputs=[
-                        model.articulation_start,
-                        model.joint_type,
-                        model.joint_parent,
-                        model.joint_child,
-                        model.joint_q_start,
-                        model.joint_qd_start,
-                        state_in.joint_q,
-                        state_in.joint_qd,
-                        model.joint_axis,
-                        model.joint_axis_start,
-                        model.joint_axis_dim,
-                        self.body_I_m,
-                        state_in.body_q,
-                        state_aug.body_q_com,
-                        model.joint_X_p,
-                        model.joint_X_c,
-                        model.gravity,
-                    ],
-                    outputs=[
-                        state_aug.joint_S_s,
-                        state_aug.body_I_s,
-                        state_aug.body_v_s,
-                        state_aug.body_f_s,
-                        state_aug.body_a_s,
-                    ],
-                    device=model.device,
-                )
-
-                if model.rigid_contact_max and (
-                    (model.ground and model.shape_ground_contact_pair_count) or model.shape_contact_pair_count
-                ):
-                    wp.launch(
-                        kernel=eval_rigid_contacts,
-                        dim=model.rigid_contact_max,
-                        inputs=[
-                            state_in.body_q,
-                            state_aug.body_v_s,
-                            model.body_com,
-                            model.shape_materials,
-                            model.shape_geo,
-                            model.shape_body,
-                            model.rigid_contact_count,
-                            model.rigid_contact_point0,
-                            model.rigid_contact_point1,
-                            model.rigid_contact_normal,
-                            model.rigid_contact_shape0,
-                            model.rigid_contact_shape1,
-                            True,
-                            self.friction_smoothing,
-                        ],
-                        outputs=[body_f],
-                        device=model.device,
-                    )
-
-                    # if model.rigid_contact_count.numpy()[0] > 0:
-                    #     print(body_f.numpy())
-
-                if model.articulation_count:
-                    # evaluate joint torques
-                    state_aug.body_ft_s.zero_()
-                    wp.launch(
-                        eval_rigid_tau,
-                        dim=model.articulation_count,
-                        inputs=[
-                            model.articulation_start,
-                            model.joint_type,
-                            model.joint_parent,
-                            model.joint_child,
-                            model.joint_q_start,
-                            model.joint_qd_start,
-                            model.joint_axis_start,
-                            model.joint_axis_dim,
-                            model.joint_axis_mode,
-                            state_in.joint_q,
-                            state_in.joint_qd,
-                            control.joint_act,
-                            model.joint_target_ke,
-                            model.joint_target_kd,
-                            model.joint_limit_lower,
-                            model.joint_limit_upper,
-                            model.joint_limit_ke,
-                            model.joint_limit_kd,
-                            state_aug.joint_S_s,
-                            state_aug.body_f_s,
-                            body_f,
-                        ],
-                        outputs=[
-                            state_aug.body_ft_s,
-                            state_aug.joint_tau,
-                        ],
-                        device=model.device,
-                    )
-
-                    # print("joint_tau:")
-                    # print(state_aug.joint_tau.numpy())
-                    # print("body_q:")
-                    # print(state_in.body_q.numpy())
-                    # print("body_qd:")
-                    # print(state_in.body_qd.numpy())
-
-                    if self._step % self.update_mass_matrix_every == 0:
-                        # build J
-                        wp.launch(
-                            eval_rigid_jacobian,
-                            dim=model.articulation_count,
-                            inputs=[
-                                model.articulation_start,
-                                self.articulation_J_start,
-                                model.joint_ancestor,
-                                model.joint_qd_start,
-                                state_aug.joint_S_s,
-                            ],
-                            outputs=[self.J],
-                            device=model.device,
-                        )
-
-                        # build M
-                        wp.launch(
-                            eval_rigid_mass,
-                            dim=model.articulation_count,
-                            inputs=[
-                                model.articulation_start,
-                                self.articulation_M_start,
-                                state_aug.body_I_s,
-                            ],
-                            outputs=[self.M],
-                            device=model.device,
-                        )
-
-                        if self.use_tile_gemm:
-                            # reshape arrays
-                            M_tiled = self.M.reshape((-1, 6 * self.joint_count, 6 * self.joint_count))
-                            J_tiled = self.J.reshape((-1, 6 * self.joint_count, self.dof_count))
-                            R_tiled = model.joint_armature.reshape((-1, self.dof_count))
-                            H_tiled = self.H.reshape((-1, self.dof_count, self.dof_count))
-                            L_tiled = self.L.reshape((-1, self.dof_count, self.dof_count))
-                            assert H_tiled.shape == (model.articulation_count, 18, 18)
-                            assert L_tiled.shape == (model.articulation_count, 18, 18)
-                            assert R_tiled.shape == (model.articulation_count, 18)
-
-                            if self.fuse_cholesky:
-                                wp.launch_tiled(
-                                    self.eval_inertia_matrix_cholesky_kernel,
-                                    dim=model.articulation_count,
-                                    inputs=[J_tiled, M_tiled, R_tiled],
-                                    outputs=[H_tiled, L_tiled],
-                                    device=model.device,
-                                    block_dim=64,
-                                )
-
-                            else:
-                                wp.launch_tiled(
-                                    self.eval_inertia_matrix_kernel,
-                                    dim=model.articulation_count,
-                                    inputs=[J_tiled, M_tiled],
-                                    outputs=[H_tiled],
-                                    device=model.device,
-                                    block_dim=256,
-                                )
-
-                                wp.launch(
-                                    eval_dense_cholesky_batched,
-                                    dim=model.articulation_count,
-                                    inputs=[
-                                        self.articulation_H_start,
-                                        self.articulation_H_rows,
-                                        self.H,
-                                        model.joint_armature,
-                                    ],
-                                    outputs=[self.L],
-                                    device=model.device,
-                                )
-
-                            # import numpy as np
-                            # J = J_tiled.numpy()
-                            # M = M_tiled.numpy()
-                            # R = R_tiled.numpy()
-                            # for i in range(model.articulation_count):
-                            #     r = R[i,:,0]
-                            #     H = J[i].T @ M[i] @ J[i]
-                            #     L = np.linalg.cholesky(H + np.diag(r))
-                            #     np.testing.assert_allclose(H, H_tiled.numpy()[i], rtol=1e-2, atol=1e-2)
-                            #     np.testing.assert_allclose(L, L_tiled.numpy()[i], rtol=1e-1, atol=1e-1)
-
-                        else:
-                            # form P = M*J
-                            wp.launch(
-                                eval_dense_gemm_batched,
-                                dim=model.articulation_count,
-                                inputs=[
-                                    self.articulation_M_rows,
-                                    self.articulation_J_cols,
-                                    self.articulation_J_rows,
-                                    False,
-                                    False,
-                                    self.articulation_M_start,
-                                    self.articulation_J_start,
-                                    # P start is the same as J start since it has the same dims as J
-                                    self.articulation_J_start,
-                                    self.M,
-                                    self.J,
-                                ],
-                                outputs=[self.P],
-                                device=model.device,
-                            )
-
-                            # form H = J^T*P
-                            wp.launch(
-                                eval_dense_gemm_batched,
-                                dim=model.articulation_count,
-                                inputs=[
-                                    self.articulation_J_cols,
-                                    self.articulation_J_cols,
-                                    # P rows is the same as J rows
-                                    self.articulation_J_rows,
-                                    True,
-                                    False,
-                                    self.articulation_J_start,
-                                    # P start is the same as J start since it has the same dims as J
-                                    self.articulation_J_start,
-                                    self.articulation_H_start,
-                                    self.J,
-                                    self.P,
-                                ],
-                                outputs=[self.H],
-                                device=model.device,
-                            )
-
-                            # compute decomposition
-                            wp.launch(
-                                eval_dense_cholesky_batched,
-                                dim=model.articulation_count,
-                                inputs=[
-                                    self.articulation_H_start,
-                                    self.articulation_H_rows,
-                                    self.H,
-                                    model.joint_armature,
-                                ],
-                                outputs=[self.L],
-                                device=model.device,
-                            )
-
-                        # print("joint_act:")
-                        # print(control.joint_act.numpy())
-                        # print("joint_tau:")
-                        # print(state_aug.joint_tau.numpy())
-                        # print("H:")
-                        # print(self.H.numpy())
-                        # print("L:")
-                        # print(self.L.numpy())
-
-                    # solve for qdd
-                    state_aug.joint_qdd.zero_()
-                    wp.launch(
-                        eval_dense_solve_batched,
-                        dim=model.articulation_count,
-                        inputs=[
-                            self.articulation_H_start,
-                            self.articulation_H_rows,
-                            self.articulation_dof_start,
-                            self.H,
-                            self.L,
-                            state_aug.joint_tau,
-                        ],
-                        outputs=[
-                            state_aug.joint_qdd,
-                            state_aug.joint_solve_tmp,
-                        ],
-                        device=model.device,
-                    )
-                    # print("joint_qdd:")
-                    # print(state_aug.joint_qdd.numpy())
-                    # print("\n\n")
-
-            # -------------------------------------
-            # integrate bodies
-
-            if model.joint_count:
-                wp.launch(
-                    kernel=integrate_generalized_joints,
-                    dim=model.joint_count,
-                    inputs=[
-                        model.joint_type,
-                        model.joint_q_start,
-                        model.joint_qd_start,
-                        model.joint_axis_dim,
-                        state_in.joint_q,
-                        state_in.joint_qd,
-                        state_aug.joint_qdd,
-                        dt,
-                    ],
-                    outputs=[state_out.joint_q, state_out.joint_qd],
-                    device=model.device,
-                )
-
-                # update maximal coordinates
-                eval_fk(model, state_out.joint_q, state_out.joint_qd, None, state_out)
-
-            self.integrate_particles(model, state_in, state_out, dt)
-
-            self._step += 1
-
-            return state_out