jaxsim 0.2.dev191__py3-none-any.whl → 0.2.dev366__py3-none-any.whl

jaxsim/rbda/aba.py

@@ -0,0 +1,295 @@
+import jax
+import jax.numpy as jnp
+import jaxlie
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+from jaxsim.math import Adjoint, Cross, Quaternion, StandardGravity
+
+from . import utils
+
+
+def aba(
+    model: js.model.JaxSimModel,
+    *,
+    base_position: jtp.VectorLike,
+    base_quaternion: jtp.VectorLike,
+    joint_positions: jtp.VectorLike,
+    base_linear_velocity: jtp.VectorLike,
+    base_angular_velocity: jtp.VectorLike,
+    joint_velocities: jtp.VectorLike,
+    joint_forces: jtp.VectorLike | None = None,
+    link_forces: jtp.MatrixLike | None = None,
+    standard_gravity: jtp.FloatLike = StandardGravity,
+) -> tuple[jtp.Vector, jtp.Vector]:
+    """
+    Compute forward dynamics using the Articulated Body Algorithm (ABA).
+
+    Args:
+        model: The model to consider.
+        base_position: The position of the base link.
+        base_quaternion: The quaternion of the base link.
+        joint_positions: The positions of the joints.
+        base_linear_velocity:
+            The linear velocity of the base link in inertial-fixed representation.
+        base_angular_velocity:
+            The angular velocity of the base link in inertial-fixed representation.
+        joint_velocities: The velocities of the joints.
+        joint_forces: The forces applied to the joints.
+        link_forces:
+            The forces applied to the links expressed in the world frame.
+        standard_gravity: The standard gravity constant.
+
+    Returns:
+        A tuple containing the base acceleration in inertial-fixed representation
+        and the joint accelerations that result from the applications of the given
+        joint and link forces.
+
+    Note:
+        The algorithm expects a quaternion with unit norm.
+    """
+
+    W_p_B, W_Q_B, s, W_v_WB, ṡ, _, _, τ, W_f, W_g = utils.process_inputs(
+        model=model,
+        base_position=base_position,
+        base_quaternion=base_quaternion,
+        joint_positions=joint_positions,
+        base_linear_velocity=base_linear_velocity,
+        base_angular_velocity=base_angular_velocity,
+        joint_velocities=joint_velocities,
+        base_linear_acceleration=None,
+        base_angular_acceleration=None,
+        joint_accelerations=None,
+        joint_forces=joint_forces,
+        link_forces=link_forces,
+        standard_gravity=standard_gravity,
+    )
+
+    W_g = jnp.atleast_2d(W_g).T
+    W_v_WB = jnp.atleast_2d(W_v_WB).T
+
+    # Get the 6D spatial inertia matrices of all links.
+    M = js.model.link_spatial_inertia_matrices(model=model)
+
+    # Get the parent array λ(i).
+    # Note: λ(0) must not be used, it's initialized to -1.
+    λ = model.kin_dyn_parameters.parent_array
+
+    # Compute the base transform.
+    W_H_B = jaxlie.SE3.from_rotation_and_translation(
+        rotation=jaxlie.SO3.from_quaternion_xyzw(xyzw=Quaternion.to_xyzw(wxyz=W_Q_B)),
+        translation=W_p_B,
+    )
+
+    # Compute 6D transforms of the base velocity.
+    W_X_B = W_H_B.adjoint()
+    B_X_W = W_H_B.inverse().adjoint()
+
+    # Compute the parent-to-child adjoints and the motion subspaces of the joints.
+    # These transforms define the relative kinematics of the entire model, including
+    # the base transform for both floating-base and fixed-base models.
+    i_X_λi, S = model.kin_dyn_parameters.joint_transforms_and_motion_subspaces(
+        joint_positions=s, base_transform=W_H_B.as_matrix()
+    )
+
+    # Allocate buffers.
+    v = jnp.zeros(shape=(model.number_of_links(), 6, 1))
+    c = jnp.zeros(shape=(model.number_of_links(), 6, 1))
+    pA = jnp.zeros(shape=(model.number_of_links(), 6, 1))
+    MA = jnp.zeros(shape=(model.number_of_links(), 6, 6))
+
+    # Allocate the buffer of transforms link -> base.
+    i_X_0 = jnp.zeros(shape=(model.number_of_links(), 6, 6))
+    i_X_0 = i_X_0.at[0].set(jnp.eye(6))
+
+    # Initialize base quantities
+    if model.floating_base():
+
+        # Base velocity v₀ in body-fixed representation.
+        v_0 = B_X_W @ W_v_WB
+        v = v.at[0].set(v_0)
+
+        # Initialize the articulated-body inertia (Mᴬ) of base link.
+        MA_0 = M[0]
+        MA = MA.at[0].set(MA_0)
+
+        # Initialize the articulated-body bias force (pᴬ) of the base link.
+        pA_0 = Cross.vx_star(v[0]) @ MA[0] @ v[0] - W_X_B.T @ jnp.vstack(W_f[0])
+        pA = pA.at[0].set(pA_0)
+
+    # ======
+    # Pass 1
+    # ======
+
+    Pass1Carry = tuple[
+        jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax
+    ]
+
+    pass_1_carry: Pass1Carry = (v, c, MA, pA, i_X_0)
+
+    # Propagate kinematics and initialize AB inertia and AB bias forces.
+    def loop_body_pass1(carry: Pass1Carry, i: jtp.Int) -> tuple[Pass1Carry, None]:
+
+        ii = i - 1
+        v, c, MA, pA, i_X_0 = carry
+
+        # Project the joint velocity into its motion subspace.
+        vJ = S[i] * ṡ[ii]
+
+        # Propagate the link velocity.
+        v_i = i_X_λi[i] @ v[λ[i]] + vJ
+        v = v.at[i].set(v_i)
+
+        c_i = Cross.vx(v[i]) @ vJ
+        c = c.at[i].set(c_i)
+
+        # Initialize the articulated-body inertia.
+        MA_i = jnp.array(M[i])
+        MA = MA.at[i].set(MA_i)
+
+        # Compute the link-to-base transform.
+        i_Xi_0 = i_X_λi[i] @ i_X_0[λ[i]]
+        i_X_0 = i_X_0.at[i].set(i_Xi_0)
+
+        # Compute link-to-world transform for the 6D force.
+        i_Xf_W = Adjoint.inverse(i_X_0[i] @ B_X_W).T
+
+        # Initialize articulated-body bias force.
+        pA_i = Cross.vx_star(v[i]) @ M[i] @ v[i] - i_Xf_W @ jnp.vstack(W_f[i])
+        pA = pA.at[i].set(pA_i)
+
+        return (v, c, MA, pA, i_X_0), None
+
+    (v, c, MA, pA, i_X_0), _ = (
+        jax.lax.scan(
+            f=loop_body_pass1,
+            init=pass_1_carry,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(v, c, MA, pA, i_X_0), None]
+    )
+
+    # ======
+    # Pass 2
+    # ======
+
+    U = jnp.zeros_like(S)
+    d = jnp.zeros(shape=(model.number_of_links(), 1))
+    u = jnp.zeros(shape=(model.number_of_links(), 1))
+
+    Pass2Carry = tuple[
+        jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax
+    ]
+
+    pass_2_carry: Pass2Carry = (U, d, u, MA, pA)
+
+    def loop_body_pass2(carry: Pass2Carry, i: jtp.Int) -> tuple[Pass2Carry, None]:
+
+        ii = i - 1
+        U, d, u, MA, pA = carry
+
+        U_i = MA[i] @ S[i]
+        U = U.at[i].set(U_i)
+
+        d_i = S[i].T @ U[i]
+        d = d.at[i].set(d_i.squeeze())
+
+        u_i = τ[ii] - S[i].T @ pA[i]
+        u = u.at[i].set(u_i.squeeze())
+
+        # Compute the articulated-body inertia and bias force of this link.
+        Ma = MA[i] - U[i] / d[i] @ U[i].T
+        pa = pA[i] + Ma @ c[i] + U[i] * (u[i] / d[i])
+
+        # Propagate them to the parent, handling the base link.
+        def propagate(
+            MA_pA: tuple[jtp.MatrixJax, jtp.MatrixJax]
+        ) -> tuple[jtp.MatrixJax, jtp.MatrixJax]:
+
+            MA, pA = MA_pA
+
+            MA_λi = MA[λ[i]] + i_X_λi[i].T @ Ma @ i_X_λi[i]
+            MA = MA.at[λ[i]].set(MA_λi)
+
+            pA_λi = pA[λ[i]] + i_X_λi[i].T @ pa
+            pA = pA.at[λ[i]].set(pA_λi)
+
+            return MA, pA
+
+        MA, pA = jax.lax.cond(
+            pred=jnp.logical_or(λ[i] != 0, model.floating_base()),
+            true_fun=propagate,
+            false_fun=lambda MA_pA: MA_pA,
+            operand=(MA, pA),
+        )
+
+        return (U, d, u, MA, pA), None
+
+    (U, d, u, MA, pA), _ = (
+        jax.lax.scan(
+            f=loop_body_pass2,
+            init=pass_2_carry,
+            xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+        )
+        if model.number_of_links() > 1
+        else [(U, d, u, MA, pA), None]
+    )
+
+    # ======
+    # Pass 3
+    # ======
+
+    if model.floating_base():
+        a0 = jnp.linalg.solve(-MA[0], pA[0])
+    else:
+        a0 = -B_X_W @ W_g
+
+    s̈ = jnp.zeros_like(s)
+    a = jnp.zeros_like(v).at[0].set(a0)
+
+    Pass3Carry = tuple[jtp.MatrixJax, jtp.VectorJax]
+    pass_3_carry = (a, s̈)
+
+    def loop_body_pass3(carry: Pass3Carry, i: jtp.Int) -> tuple[Pass3Carry, None]:
+
+        ii = i - 1
+        a, s̈ = carry
+
+        # Propagate the link acceleration.
+        a_i = i_X_λi[i] @ a[λ[i]] + c[i]
+
+        # Compute the joint acceleration.
+        s̈_ii = (u[i] - U[i].T @ a_i) / d[i]
+        s̈ = s̈.at[ii].set(s̈_ii.squeeze())
+
+        # Sum the joint acceleration to the parent link acceleration.
+        a_i = a_i + S[i] * s̈[ii]
+        a = a.at[i].set(a_i)
+
+        return (a, s̈), None
+
+    (a, s̈), _ = (
+        jax.lax.scan(
+            f=loop_body_pass3,
+            init=pass_3_carry,
+            xs=jnp.arange(1, model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(a, s̈), None]
+    )
+
+    # ==============
+    # Adjust outputs
+    # ==============
+
+    # TODO: remove vstack and shape=(6, 1)?
+    if model.floating_base():
+        # Convert the base acceleration to inertial-fixed representation,
+        # and add gravity.
+        B_a_WB = a[0]
+        W_a_WB = W_X_B @ B_a_WB + W_g
+    else:
+        W_a_WB = jnp.zeros(6)
+
+    return W_a_WB.squeeze(), jnp.atleast_1d(s̈.squeeze())

jaxsim/rbda/collidable_points.py

@@ -0,0 +1,142 @@
+import jax
+import jax.numpy as jnp
+import jaxlie
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+from jaxsim.math import Adjoint, Quaternion, Skew
+
+from . import utils
+
+
+def collidable_points_pos_vel(
+    model: js.model.JaxSimModel,
+    *,
+    base_position: jtp.Vector,
+    base_quaternion: jtp.Vector,
+    joint_positions: jtp.Vector,
+    base_linear_velocity: jtp.Vector,
+    base_angular_velocity: jtp.Vector,
+    joint_velocities: jtp.Vector,
+) -> tuple[jtp.Matrix, jtp.Matrix]:
+    """
+
+    Compute the position and linear velocity of collidable points in the world frame.
+
+    Args:
+        model: The model to consider.
+        base_position: The position of the base link.
+        base_quaternion: The quaternion of the base link.
+        joint_positions: The positions of the joints.
+        base_linear_velocity:
+            The linear velocity of the base link in inertial-fixed representation.
+        base_angular_velocity:
+            The angular velocity of the base link in inertial-fixed representation.
+        joint_velocities: The velocities of the joints.
+
+    Returns:
+        A tuple containing the position and linear velocity of collidable points.
+    """
+
+    if len(model.kin_dyn_parameters.contact_parameters.body) == 0:
+        return jnp.array(0).astype(float), jnp.empty(0).astype(float)
+
+    W_p_B, W_Q_B, s, W_v_WB, ṡ, _, _, _, _, _ = utils.process_inputs(
+        model=model,
+        base_position=base_position,
+        base_quaternion=base_quaternion,
+        joint_positions=joint_positions,
+        base_linear_velocity=base_linear_velocity,
+        base_angular_velocity=base_angular_velocity,
+        joint_velocities=joint_velocities,
+    )
+
+    # Get the parent array λ(i).
+    # Note: λ(0) must not be used, it's initialized to -1.
+    λ = model.kin_dyn_parameters.parent_array
+
+    # Compute the base transform.
+    W_H_B = jaxlie.SE3.from_rotation_and_translation(
+        rotation=jaxlie.SO3.from_quaternion_xyzw(xyzw=Quaternion.to_xyzw(wxyz=W_Q_B)),
+        translation=W_p_B,
+    )
+
+    # Compute the parent-to-child adjoints and the motion subspaces of the joints.
+    # These transforms define the relative kinematics of the entire model, including
+    # the base transform for both floating-base and fixed-base models.
+    i_X_λi, S = model.kin_dyn_parameters.joint_transforms_and_motion_subspaces(
+        joint_positions=s, base_transform=W_H_B.as_matrix()
+    )
+
+    # Allocate buffer of transforms world -> link and initialize the base pose.
+    W_X_i = jnp.zeros(shape=(model.number_of_links(), 6, 6))
+    W_X_i = W_X_i.at[0].set(Adjoint.inverse(i_X_λi[0]))
+
+    # Allocate buffer of 6D inertial-fixed velocities and initialize the base velocity.
+    W_v_Wi = jnp.zeros(shape=(model.number_of_links(), 6))
+    W_v_Wi = W_v_Wi.at[0].set(W_v_WB)
+
+    # ====================
+    # Propagate kinematics
+    # ====================
+
+    PropagateTransformsCarry = tuple[jtp.MatrixJax, jtp.Matrix]
+    propagate_transforms_carry: PropagateTransformsCarry = (W_X_i, W_v_Wi)
+
+    def propagate_kinematics(
+        carry: PropagateTransformsCarry, i: jtp.Int
+    ) -> tuple[PropagateTransformsCarry, None]:
+
+        ii = i - 1
+        W_X_i, W_v_Wi = carry
+
+        # Compute the parent to child 6D transform.
+        λi_X_i = Adjoint.inverse(adjoint=i_X_λi[i])
+
+        # Compute the world to child 6D transform.
+        W_Xi_i = W_X_i[λ[i]] @ λi_X_i
+        W_X_i = W_X_i.at[i].set(W_Xi_i)
+
+        # Propagate the 6D velocity
+        W_vi_Wi = W_v_Wi[λ[i]] + W_X_i[i] @ (S[i] * ṡ[ii]).squeeze()
+        W_v_Wi = W_v_Wi.at[i].set(W_vi_Wi)
+
+        return (W_X_i, W_v_Wi), None
+
+    (W_X_i, W_v_Wi), _ = (
+        jax.lax.scan(
+            f=propagate_kinematics,
+            init=propagate_transforms_carry,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(W_X_i, W_v_Wi), None]
+    )
+
+    # ==================================================
+    # Compute position and velocity of collidable points
+    # ==================================================
+
+    def process_point_kinematics(
+        Li_p_C: jtp.VectorJax, parent_body: jtp.Int
+    ) -> tuple[jtp.VectorJax, jtp.VectorJax]:
+        # Compute the position of the collidable point
+        W_p_Ci = (
+            Adjoint.to_transform(adjoint=W_X_i[parent_body]) @ jnp.hstack([Li_p_C, 1])
+        )[0:3]
+
+        # Compute the linear part of the mixed velocity Ci[W]_v_{W,Ci}
+        CW_vl_WCi = (
+            jnp.block([jnp.eye(3), -Skew.wedge(vector=W_p_Ci).squeeze()])
+            @ W_v_Wi[parent_body].squeeze()
+        )
+
+        return W_p_Ci, CW_vl_WCi
+
+    # Process all the collidable points in parallel
+    W_p_Ci, CW_vl_WC = jax.vmap(process_point_kinematics)(
+        model.kin_dyn_parameters.contact_parameters.point,
+        jnp.array(model.kin_dyn_parameters.contact_parameters.body),
+    )
+
+    return W_p_Ci, CW_vl_WC

jaxsim/{physics/algos → rbda}/crba.py

@@ -1,68 +1,68 @@
-from typing import Tuple
-
 import jax
 import jax.numpy as jnp
-import numpy as np
 
+import jaxsim.api as js
 import jaxsim.typing as jtp
-from jaxsim.physics.model.physics_model import PhysicsModel
 
 from . import utils
 
 
-def crba(model: PhysicsModel, q: jtp.Vector) -> jtp.Matrix:
+def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Matrix:
     """
-    Compute the Composite Rigid-Body Inertia Matrix (CRBA) for an articulated body or robot given joint positions.
+    Compute the free-floating mass matrix using the Composite Rigid-Body Algorithm (CRBA).
 
     Args:
-        model (PhysicsModel): The physics model of the articulated body or robot.
-        q (jtp.Vector): Joint positions (Generalized coordinates).
+        model: The model to consider.
+        joint_positions: The positions of the joints.
 
     Returns:
-        jtp.Matrix: The Composite Rigid-Body Inertia Matrix (CRBA) of the articulated body or robot.
+        The free-floating mass matrix of the model in body-fixed representation.
     """
 
-    _, q, _, _, _, _ = utils.process_inputs(
-        physics_model=model, xfb=None, q=q, qd=None, tau=None, f_ext=None
+    _, _, s, _, _, _, _, _, _, _ = utils.process_inputs(
+        model=model, joint_positions=joint_positions
     )
 
-    Xtree = model.tree_transforms
-    Mc = model.spatial_inertias
-    S = model.motion_subspaces(q=q)
-    Xj = model.joint_transforms(q=q)
+    # Get the 6D spatial inertia matrices of all links.
+    Mc = js.model.link_spatial_inertia_matrices(model=model)
 
-    Xup = jnp.zeros_like(Xtree)
-    i_X_0 = jnp.zeros_like(Xtree)
-    i_X_0 = i_X_0.at[0].set(jnp.eye(6))
+    # Get the parent array λ(i).
+    # Note: λ(0) must not be used, it's initialized to -1.
+    λ = model.kin_dyn_parameters.parent_array
+
+    # Compute the parent-to-child adjoints and the motion subspaces of the joints.
+    # These transforms define the relative kinematics of the entire model, including
+    # the base transform for both floating-base and fixed-base models.
+    i_X_λi, S = model.kin_dyn_parameters.joint_transforms_and_motion_subspaces(
+        joint_positions=s, base_transform=jnp.eye(4)
+    )
 
-    # Parent array mapping: i -> λ(i).
-    # Exception: λ(0) must not be used, it's initialized to -1.
-    λ = model.parent
+    # Allocate the buffer of transforms link -> base.
+    i_X_0 = jnp.zeros(shape=(model.number_of_links(), 6, 6))
+    i_X_0 = i_X_0.at[0].set(jnp.eye(6))
 
     # ====================
     # Propagate kinematics
     # ====================
 
-    ForwardPassCarry = Tuple[jtp.MatrixJax, jtp.MatrixJax]
-    forward_pass_carry = (Xup, i_X_0)
+    ForwardPassCarry = tuple[jtp.MatrixJax]
+    forward_pass_carry: ForwardPassCarry = (i_X_0,)
 
     def propagate_kinematics(
         carry: ForwardPassCarry, i: jtp.Int
-    ) -> Tuple[ForwardPassCarry, None]:
-        Xup, i_X_0 = carry
+    ) -> tuple[ForwardPassCarry, None]:
 
-        Xup_i = Xj[i] @ Xtree[i]
-        Xup = Xup.at[i].set(Xup_i)
+        (i_X_0,) = carry
 
-        i_X_0_i = Xup[i] @ i_X_0[λ[i]]
+        i_X_0_i = i_X_λi[i] @ i_X_0[λ[i]]
         i_X_0 = i_X_0.at[i].set(i_X_0_i)
 
-        return (Xup, i_X_0), None
+        return (i_X_0,), None
 
-    (Xup, i_X_0), _ = jax.lax.scan(
+    (i_X_0,), _ = jax.lax.scan(
         f=propagate_kinematics,
         init=forward_pass_carry,
-        xs=np.arange(start=1, stop=model.NB),
+        xs=jnp.arange(start=1, stop=model.number_of_links()),
     )
 
     # ===================
@@ -71,16 +71,17 @@ def crba(model: PhysicsModel, q: jtp.Vector) -> jtp.Matrix:
 
     M = jnp.zeros(shape=(6 + model.dofs(), 6 + model.dofs()))
 
-    BackwardPassCarry = Tuple[jtp.MatrixJax, jtp.MatrixJax]
-    backward_pass_carry = (Mc, M)
+    BackwardPassCarry = tuple[jtp.MatrixJax, jtp.MatrixJax]
+    backward_pass_carry: BackwardPassCarry = (Mc, M)
 
     def backward_pass(
         carry: BackwardPassCarry, i: jtp.Int
-    ) -> Tuple[BackwardPassCarry, None]:
+    ) -> tuple[BackwardPassCarry, None]:
+
         ii = i - 1
         Mc, M = carry
 
-        Mc_λi = Mc[λ[i]] + Xup[i].T @ Mc[i] @ Xup[i]
+        Mc_λi = Mc[λ[i]] + i_X_λi[i].T @ Mc[i] @ i_X_λi[i]
         Mc = Mc.at[λ[i]].set(Mc_λi)
 
         Fi = Mc[i] @ S[i]
@@ -89,13 +90,13 @@ def crba(model: PhysicsModel, q: jtp.Vector) -> jtp.Matrix:
 
         j = i
 
-        CarryInnerFn = Tuple[jtp.Int, jtp.MatrixJax, jtp.MatrixJax]
+        CarryInnerFn = tuple[jtp.Int, jtp.MatrixJax, jtp.MatrixJax]
         carry_inner_fn = (j, Fi, M)
 
         def while_loop_body(carry: CarryInnerFn) -> CarryInnerFn:
             j, Fi, M = carry
 
-            Fi = Xup[j].T @ Fi
+            Fi = i_X_λi[j].T @ Fi
             j = λ[j]
             jj = j - 1
 
@@ -108,8 +109,8 @@ def crba(model: PhysicsModel, q: jtp.Vector) -> jtp.Matrix:
 
         # The following functions are part of a (rather messy) workaround for computing
         # a while loop using a for loop with fixed number of iterations.
-        def inner_fn(carry: CarryInnerFn, k: jtp.Int) -> Tuple[CarryInnerFn, None]:
-            def compute_inner(carry: CarryInnerFn) -> Tuple[CarryInnerFn, None]:
+        def inner_fn(carry: CarryInnerFn, k: jtp.Int) -> tuple[CarryInnerFn, None]:
+            def compute_inner(carry: CarryInnerFn) -> tuple[CarryInnerFn, None]:
                 j, Fi, M = carry
                 out = jax.lax.cond(
                     pred=(λ[j] > 0),
@@ -130,7 +131,7 @@ def crba(model: PhysicsModel, q: jtp.Vector) -> jtp.Matrix:
         (j, Fi, M), _ = jax.lax.scan(
             f=inner_fn,
             init=carry_inner_fn,
-            xs=np.flip(np.arange(start=1, stop=model.NB)),
+            xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
         )
 
         Fi = i_X_0[j].T @ Fi
@@ -145,10 +146,10 @@ def crba(model: PhysicsModel, q: jtp.Vector) -> jtp.Matrix:
     (Mc, M), _ = jax.lax.scan(
         f=backward_pass,
         init=backward_pass_carry,
-        xs=np.flip(np.arange(start=1, stop=model.NB)),
+        xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
     )
 
-    # Store the locked 6D rigid-body inertia matrix Mbb ∈ ℝ⁶ˣ⁶
+    # Store the locked 6D rigid-body inertia matrix Mbb ∈ ℝ⁶ˣ⁶.
     M = M.at[0:6, 0:6].set(Mc[0])
 
     return M