jaxsim 0.6.1.dev13__py3-none-any.whl → 0.6.2.dev102__py3-none-any.whl

jaxsim/api/ode.py

@@ -1,230 +1,24 @@
-from typing import Any, Protocol
-
 import jax
 import jax.numpy as jnp
 
 import jaxsim.api as js
-import jaxsim.rbda
 import jaxsim.typing as jtp
-from jaxsim.integrators import Time
-from jaxsim.math import Quaternion
-from jaxsim.rbda import contacts
+from jaxsim.api.data import JaxSimModelData
+from jaxsim.math import Quaternion, Skew
 
 from .common import VelRepr
-from .ode_data import ODEState
-
-
-class SystemDynamicsFromModelAndData(Protocol):
-    """
-    Protocol defining the signature of a function computing the system dynamics
-    given a model and data object.
-    """
-
-    def __call__(
-        self,
-        model: js.model.JaxSimModel,
-        data: js.data.JaxSimModelData,
-        **kwargs: dict[str, Any],
-    ) -> tuple[ODEState, dict[str, Any]]:
-        """
-        Compute the system dynamics given a model and data object.
-
-        Args:
-            model: The model to consider.
-            data: The data of the considered model.
-            **kwargs: Additional keyword arguments.
-
-        Returns:
-            A tuple with an `ODEState` object storing in each of its attributes the
-            corresponding derivative, and the dictionary of auxiliary data returned
-            by the system dynamics evaluation.
-        """
-
-        pass
-
-
-def wrap_system_dynamics_for_integration(
-    *,
-    system_dynamics: SystemDynamicsFromModelAndData,
-    **kwargs: dict[str, Any],
-) -> jaxsim.integrators.common.SystemDynamics[ODEState, ODEState]:
-    """
-    Wrap the system dynamics considered by JaxSim integrators in a generic
-    `f(x, t, **u, **parameters)` function.
-
-    Args:
-        system_dynamics: The system dynamics to wrap.
-        **kwargs: Additional kwargs to close over the system dynamics.
-
-    Returns:
-        The system dynamics closed over the additional kwargs to be used by
-        JaxSim integrators.
-    """
-
-    # Close `system_dynamics` over additional kwargs.
-    # Similarly to what done in `jaxsim.api.model.step`, to be future-proof, we use the
-    # following logic to allow the caller to close over arguments having the same name
-    # of the ones used in the `wrap_system_dynamics_for_integration` function.
-    kwargs = kwargs.copy() if kwargs is not None else {}
-    colliding_system_dynamics_kwargs = kwargs.pop("system_dynamics_kwargs", {})
-    system_dynamics_kwargs = kwargs | colliding_system_dynamics_kwargs
-
-    # Remove `model` and `data` for backward compatibility.
-    # It's no longer necessary to close over them at this stage, as this is always
-    # done in `jaxsim.api.model.step`.
-    # We can remove the following lines in a few releases.
-    _ = system_dynamics_kwargs.pop("data", None)
-    _ = system_dynamics_kwargs.pop("model", None)
-
-    # Create the function with the signature expected by our generic integrators.
-    # Note that our system dynamics is time independent.
-    def f(x: ODEState, t: Time, **kwargs_f) -> tuple[ODEState, dict[str, Any]]:
-
-        # Get the data and model objects from the kwargs.
-        data_f = kwargs_f.pop("data")
-        model_f = kwargs_f.pop("model")
-
-        # Update the state and time stored inside data.
-        with data_f.editable(validate=True) as data_rw:
-            data_rw.state = x
-
-        # Evaluate the system dynamics, allowing to override the kwargs originally
-        # passed when the closure was created.
-        return system_dynamics(
-            model=model_f,
-            data=data_rw,
-            **(system_dynamics_kwargs | kwargs_f),
-        )
-
-    f: jaxsim.integrators.common.SystemDynamics[ODEState, ODEState]
-    return f
-
 
 # ==================================
 # Functions defining system dynamics
 # ==================================
 
 
-@jax.jit
-@js.common.named_scope
-def system_velocity_dynamics(
-    model: js.model.JaxSimModel,
-    data: js.data.JaxSimModelData,
-    *,
-    link_forces: jtp.Vector | None = None,
-    joint_force_references: jtp.Vector | None = None,
-) -> tuple[jtp.Vector, jtp.Vector, dict[str, Any]]:
-    """
-    Compute the dynamics of the system velocity.
-
-    Args:
-        model: The model to consider.
-        data: The data of the considered model.
-        link_forces:
-            The 6D forces to apply to the links expressed in the frame corresponding to
-            the velocity representation of `data`.
-        joint_force_references: The joint force references to apply.
-
-    Returns:
-        A tuple containing the derivative of the base 6D velocity in inertial-fixed
-        representation, the derivative of the joint velocities, and auxiliary data
-        returned by the system dynamics evaluation.
-    """
-
-    # Build link forces if not provided.
-    # These forces are expressed in the frame corresponding to the velocity
-    # representation of data.
-    O_f_L = (
-        jnp.atleast_2d(link_forces.squeeze())
-        if link_forces is not None
-        else jnp.zeros((model.number_of_links(), 6))
-    ).astype(float)
-
-    # We expect that the 6D forces included in the `link_forces` argument are expressed
-    # in the frame corresponding to the velocity representation of `data`.
-    references = js.references.JaxSimModelReferences.build(
-        model=model,
-        link_forces=O_f_L,
-        joint_force_references=joint_force_references,
-        data=data,
-        velocity_representation=data.velocity_representation,
-    )
-
-    # ======================
-    # Compute contact forces
-    # ======================
-
-    # Initialize the 6D forces W_f ∈ ℝ^{n_L × 6} applied to links due to contact
-    # with the terrain.
-    W_f_L_terrain = jnp.zeros_like(O_f_L).astype(float)
-
-    # Initialize a dictionary of auxiliary data.
-    # This dictionary is used to store additional data computed by the contact model.
-    aux_data = {}
-
-    if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
-
-        with (
-            data.switch_velocity_representation(VelRepr.Inertial),
-            references.switch_velocity_representation(VelRepr.Inertial),
-        ):
-
-            # Compute the 6D forces W_f ∈ ℝ^{n_c × 6} applied to each collidable point
-            # along with contact-specific auxiliary states.
-            W_f_C, aux_data = js.contact.collidable_point_dynamics(
-                model=model,
-                data=data,
-                link_forces=references.link_forces(model=model, data=data),
-                joint_force_references=references.joint_force_references(model=model),
-            )
-
-            # Compute the 6D forces applied to the links equivalent to the forces applied
-            # to the frames associated to the collidable points.
-            W_f_L_terrain = model.contact_model.link_forces_from_contact_forces(
-                model=model,
-                data=data,
-                contact_forces=W_f_C,
-            )
-
-    # ===========================
-    # Compute system acceleration
-    # ===========================
-
-    # Compute the total link forces.
-    with (
-        data.switch_velocity_representation(VelRepr.Inertial),
-        references.switch_velocity_representation(VelRepr.Inertial),
-    ):
-
-        # Sum the contact forces just computed with the link forces applied by the user.
-        references = references.apply_link_forces(
-            model=model,
-            data=data,
-            forces=W_f_L_terrain,
-            additive=True,
-        )
-
-        # Get the link forces in inertial-fixed representation.
-        f_L_total = references.link_forces(model=model, data=data)
-
-        # Compute the system acceleration in inertial-fixed representation.
-        # This representation is useful for integration purpose.
-        W_v̇_WB, s̈ = system_acceleration(
-            model=model,
-            data=data,
-            joint_force_references=joint_force_references,
-            link_forces=f_L_total,
-        )
-
-    return W_v̇_WB, s̈, aux_data
-
-
 def system_acceleration(
     model: js.model.JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
     link_forces: jtp.MatrixLike | None = None,
-    joint_force_references: jtp.VectorLike | None = None,
+    joint_torques: jtp.VectorLike | None = None,
 ) -> tuple[jtp.Vector, jtp.Vector]:
     """
     Compute the system acceleration in the active representation.
@@ -235,7 +29,7 @@ def system_acceleration(
         link_forces:
             The 6D forces to apply to the links expressed in the same
             velocity representation of data.
-        joint_force_references: The joint force references to apply.
+        joint_torques: The joint torques applied to the joints.
 
     Returns:
         A tuple containing the base 6D acceleration in the active representation
@@ -253,80 +47,6 @@ def system_acceleration(
         else jnp.zeros((model.number_of_links(), 6))
     ).astype(float)
 
-    # Build joint torques if not provided.
-    τ_references = (
-        jnp.atleast_1d(joint_force_references.squeeze())
-        if joint_force_references is not None
-        else jnp.zeros_like(data.joint_positions())
-    ).astype(float)
-
-    # ====================
-    # Enforce joint limits
-    # ====================
-
-    τ_position_limit = jnp.zeros_like(τ_references).astype(float)
-
-    if model.dofs() > 0:
-
-        # Stiffness and damper parameters for the joint position limits.
-        k_j = jnp.array(
-            model.kin_dyn_parameters.joint_parameters.position_limit_spring
-        ).astype(float)
-        d_j = jnp.array(
-            model.kin_dyn_parameters.joint_parameters.position_limit_damper
-        ).astype(float)
-
-        # Compute the joint position limit violations.
-        lower_violation = jnp.clip(
-            data.state.physics_model.joint_positions
-            - model.kin_dyn_parameters.joint_parameters.position_limits_min,
-            max=0.0,
-        )
-
-        upper_violation = jnp.clip(
-            data.state.physics_model.joint_positions
-            - model.kin_dyn_parameters.joint_parameters.position_limits_max,
-            min=0.0,
-        )
-
-        # Compute the joint position limit torque.
-        τ_position_limit -= jnp.diag(k_j) @ (lower_violation + upper_violation)
-
-        τ_position_limit -= (
-            jnp.positive(τ_position_limit)
-            * jnp.diag(d_j)
-            @ data.state.physics_model.joint_velocities
-        )
-
-    # ====================
-    # Joint friction model
-    # ====================
-
-    τ_friction = jnp.zeros_like(τ_references).astype(float)
-
-    if model.dofs() > 0:
-
-        # Static and viscous joint friction parameters
-        kc = jnp.array(
-            model.kin_dyn_parameters.joint_parameters.friction_static
-        ).astype(float)
-        kv = jnp.array(
-            model.kin_dyn_parameters.joint_parameters.friction_viscous
-        ).astype(float)
-
-        # Compute the joint friction torque.
-        τ_friction = -(
-            jnp.diag(kc) @ jnp.sign(data.state.physics_model.joint_velocities)
-            + jnp.diag(kv) @ data.state.physics_model.joint_velocities
-        )
-
-    # ========================
-    # Compute forward dynamics
-    # ========================
-
-    # Compute the total joint forces.
-    τ_total = τ_references + τ_friction + τ_position_limit
-
     # Store the link forces in a references object.
     references = js.references.JaxSimModelReferences.build(
         model=model,
@@ -345,7 +65,7 @@ def system_acceleration(
     v̇_WB, s̈ = js.model.forward_dynamics_aba(
         model=model,
         data=data,
-        joint_forces=τ_total,
+        joint_forces=joint_torques,
         link_forces=references.link_forces(model=model, data=data),
     )
 
@@ -359,7 +79,7 @@ def system_position_dynamics(
     data: js.data.JaxSimModelData,
     baumgarte_quaternion_regularization: jtp.FloatLike = 1.0,
 ) -> tuple[jtp.Vector, jtp.Vector, jtp.Vector]:
-    """
+    r"""
     Compute the dynamics of the system position.
 
     Args:
@@ -371,16 +91,18 @@ def system_position_dynamics(
     Returns:
         A tuple containing the derivative of the base position, the derivative of the
         base quaternion, and the derivative of the joint positions.
-    """
 
-    ṡ = data.joint_velocities(model=model)
-    W_Q_B = data.base_orientation(dcm=False)
-
-    with data.switch_velocity_representation(velocity_representation=VelRepr.Mixed):
-        W_ṗ_B = data.base_velocity()[0:3]
+    Note:
+        In inertial-fixed representation, the linear component of the base velocity is not
+        the derivative of the base position. In fact, the base velocity is defined as:
+        :math:`{} ^W v_{W, B} = \begin{bmatrix} {} ^W \dot{p}_B S({} ^W \omega_{W, B}) {} ^W p _B\\ {} ^W \omega_{W, B} \end{bmatrix}`.
+        Where :math:`S(\cdot)` is the skew-symmetric matrix operator.
+    """
 
-    with data.switch_velocity_representation(velocity_representation=VelRepr.Inertial):
-        W_ω_WB = data.base_velocity()[3:6]
+    ṡ = data.joint_velocities
+    W_Q_B = data.base_orientation
+    W_ω_WB = data.base_velocity[3:6]
+    W_ṗ_B = data.base_velocity[0:3] + Skew.wedge(W_ω_WB) @ data.base_position
 
     W_Q̇_B = Quaternion.derivative(
         quaternion=W_Q_B,
@@ -399,9 +121,9 @@ def system_dynamics(
     data: js.data.JaxSimModelData,
     *,
     link_forces: jtp.Vector | None = None,
-    joint_force_references: jtp.Vector | None = None,
+    joint_torques: jtp.Vector | None = None,
     baumgarte_quaternion_regularization: jtp.FloatLike = 1.0,
-) -> tuple[ODEState, dict[str, Any]]:
+) -> JaxSimModelData:
     """
     Compute the dynamics of the system.
 
@@ -411,57 +133,32 @@ def system_dynamics(
         link_forces:
             The 6D forces to apply to the links expressed in the frame corresponding to
             the velocity representation of `data`.
-        joint_force_references: The joint force references to apply.
+        joint_torques: The joint torques acting on the joints.
         baumgarte_quaternion_regularization:
             The Baumgarte regularization coefficient used to adjust the norm of the
             quaternion (only used in integrators not operating on the SO(3) manifold).
 
     Returns:
-        A tuple with an `ODEState` object storing in each of its attributes the
+        A tuple with an `JaxSimModelData` object storing in each of its attributes the
         corresponding derivative, and the dictionary of auxiliary data returned
         by the system dynamics evaluation.
     """
 
-    # Compute the accelerations and the material deformation rate.
-    W_v̇_WB, s̈, aux_dict = system_velocity_dynamics(
-        model=model,
-        data=data,
-        joint_force_references=joint_force_references,
-        link_forces=link_forces,
-    )
-
-    # Initialize the dictionary storing the derivative of the additional state variables
-    # that extend the state vector of the integrated ODE system.
-    extended_ode_state = {}
-
-    match model.contact_model:
-
-        case contacts.SoftContacts():
-            extended_ode_state["tangential_deformation"] = aux_dict["m_dot"]
-
-        case contacts.ViscoElasticContacts():
-
-            extended_ode_state["tangential_deformation"] = jnp.zeros_like(
-                data.state.extended["tangential_deformation"]
-            )
-
-        case contacts.RigidContacts() | contacts.RelaxedRigidContacts():
-            pass
-
-        case _:
-            raise ValueError(f"Invalid contact model: {model.contact_model}")
+    with data.switch_velocity_representation(velocity_representation=VelRepr.Inertial):
+        W_v̇_WB, s̈ = system_acceleration(
+            model=model,
+            data=data,
+            joint_torques=joint_torques,
+            link_forces=link_forces,
+        )
 
-    # Extract the velocities.
-    W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(
-        model=model,
-        data=data,
-        baumgarte_quaternion_regularization=baumgarte_quaternion_regularization,
-    )
+        W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(
+            model=model,
+            data=data,
+            baumgarte_quaternion_regularization=baumgarte_quaternion_regularization,
+        )
 
-    # Create an ODEState object populated with the derivative of each leaf.
-    # Our integrators, operating on generic pytrees, will be able to handle it
-    # automatically as state derivative.
-    ode_state_derivative = ODEState.build_from_jaxsim_model(
+    ode_state_derivative = JaxSimModelData.build(
         model=model,
         base_position=W_ṗ_B,
         base_quaternion=W_Q̇_B,
@@ -469,7 +166,6 @@ def system_dynamics(
         base_linear_velocity=W_v̇_WB[0:3],
         base_angular_velocity=W_v̇_WB[3:6],
         joint_velocities=s̈,
-        **extended_ode_state,
     )
 
-    return ode_state_derivative, aux_dict
+    return ode_state_derivative

jaxsim/api/references.py

@@ -242,7 +242,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             )(W_f_L, W_H_L)
 
         # The f_L output is either L_f_L or LW_f_L, depending on the representation.
-        W_H_L = js.model.forward_kinematics(model=model, data=data)
+        W_H_L = data._link_transforms
         f_L = convert(W_f_L=W_f_L[link_idxs, :], W_H_L=W_H_L[link_idxs, :, :])
 
         return f_L
@@ -450,7 +450,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             )(f_L, W_H_L)
 
         # The f_L input is either L_f_L or LW_f_L, depending on the representation.
-        W_H_L = js.model.forward_kinematics(model=model, data=data)
+        W_H_L = data._link_transforms
         W_f_L = convert_using_link_frame(f_L=f_L, W_H_L=W_H_L[link_idxs, :, :])
 
         return replace(forces=self._link_forces.at[link_idxs, :].set(W_f0_L + W_f_L))

jaxsim/exceptions.py

@@ -23,8 +23,8 @@ def raise_if(
 
     # Disable host callback if running on unsupported hardware or if the user
     # explicitly disabled it.
-    if jax.devices()[0].platform in {"tpu", "METAL"} or os.environ.get(
-        "JAXSIM_DISABLE_EXCEPTIONS", 0
+    if jax.devices()[0].platform in {"tpu", "METAL"} or not os.environ.get(
+        "JAXSIM_ENABLE_EXCEPTIONS", 0
     ):
         return
 

jaxsim/math/__init__.py

@@ -1,6 +1,3 @@
-# Define the default standard gravity constant.
-StandardGravity = 9.81
-
 from .adjoint import Adjoint
 from .cross import Cross
 from .inertia import Inertia
@@ -11,3 +8,7 @@ from .transform import Transform
 from .utils import safe_norm
 
 from .joint_model import JointModel, supported_joint_motion  # isort:skip
+
+
+# Define the default standard gravity constant.
+STANDARD_GRAVITY = -9.81

jaxsim/math/joint_model.py

@@ -7,12 +7,10 @@ import jaxlie
 from jax_dataclasses import Static
 
 import jaxsim.typing as jtp
+from jaxsim.math import Rotation
 from jaxsim.parsers.descriptions import JointGenericAxis, JointType, ModelDescription
 from jaxsim.parsers.kinematic_graph import KinematicGraphTransforms
 
-from .rotation import Rotation
-from .transform import Transform
-
 
 @jax_dataclasses.pytree_dataclass
 class JointModel:
@@ -113,60 +111,6 @@ class JointModel:
             joint_axis=tuple(JointGenericAxis(axis=j.axis) for j in ordered_joints),
         )
 
-    def parent_H_child(
-        self, joint_index: jtp.IntLike, joint_position: jtp.VectorLike
-    ) -> tuple[jtp.Matrix, jtp.Array]:
-        r"""
-        Compute the homogeneous transformation between the parent link and
-        the child link of a joint, and the corresponding motion subspace.
-
-        Args:
-            joint_index: The index of the joint.
-            joint_position: The position of the joint.
-
-        Returns:
-            A tuple containing the homogeneous transformation
-            :math:`{}^{\lambda(i)} \mathbf{H}_i(s)`
-            and the motion subspace :math:`\mathbf{S}(s)`.
-        """
-
-        i = joint_index
-        s = joint_position
-
-        # Get the components of the joint model.
-        λ_Hi_pre = self.parent_H_predecessor(joint_index=i)
-        pre_Hi_suc, S = self.predecessor_H_successor(joint_index=i, joint_position=s)
-        suc_Hi_i = self.successor_H_child(joint_index=i)
-
-        # Compose all the transforms.
-        return λ_Hi_pre @ pre_Hi_suc @ suc_Hi_i, S
-
-    @jax.jit
-    def child_H_parent(
-        self, joint_index: jtp.IntLike, joint_position: jtp.VectorLike
-    ) -> tuple[jtp.Matrix, jtp.Array]:
-        r"""
-        Compute the homogeneous transformation between the child link and
-        the parent link of a joint, and the corresponding motion subspace.
-
-        Args:
-            joint_index: The index of the joint.
-            joint_position: The position of the joint.
-
-        Returns:
-            A tuple containing the homogeneous transformation
-            :math:`{}^{i} \mathbf{H}_{\lambda(i)}(s)`
-            and the motion subspace :math:`\mathbf{S}(s)`.
-        """
-
-        λ_Hi_i, S = self.parent_H_child(
-            joint_index=joint_index, joint_position=joint_position
-        )
-
-        i_Hi_λ = Transform.inverse(λ_Hi_i)
-
-        return i_Hi_λ, S
-
     def parent_H_predecessor(self, joint_index: jtp.IntLike) -> jtp.Matrix:
         r"""
         Return the homogeneous transformation between the parent link and
@@ -182,31 +126,6 @@ class JointModel:
 
         return self.λ_H_pre[joint_index]
 
-    def predecessor_H_successor(
-        self, joint_index: jtp.IntLike, joint_position: jtp.VectorLike
-    ) -> tuple[jtp.Matrix, jtp.Array]:
-        r"""
-        Compute the homogeneous transformation between the predecessor and
-        the successor frame of a joint, and the corresponding motion subspace.
-
-        Args:
-            joint_index: The index of the joint.
-            joint_position: The position of the joint.
-
-        Returns:
-            A tuple containing the homogeneous transformation
-            :math:`{}^{\text{pre}(i)} \mathbf{H}_{\text{suc}(i)}(s)`
-            and the motion subspace :math:`\mathbf{S}(s)`.
-        """
-
-        pre_H_suc, S = supported_joint_motion(
-            self.joint_types[joint_index],
-            joint_position,
-            self.joint_axis[joint_index].axis,
-        )
-
-        return pre_H_suc, S
-
     def successor_H_child(self, joint_index: jtp.IntLike) -> jtp.Matrix:
         r"""
         Return the homogeneous transformation between the successor frame and
@@ -225,65 +144,56 @@ class JointModel:
 
 @jax.jit
 def supported_joint_motion(
-    joint_type: jtp.IntLike,
-    joint_position: jtp.VectorLike,
-    joint_axis: jtp.VectorLike | None = None,
-    /,
-) -> tuple[jtp.Matrix, jtp.Array]:
+    joint_types: jtp.Array, joint_positions: jtp.Matrix, joint_axes: jtp.Matrix
+) -> jtp.Matrix:
     """
-    Compute the homogeneous transformation and motion subspace of a joint.
+    Compute the transforms of the joints.
 
     Args:
-        joint_type: The type of the joint.
-        joint_position: The position of the joint.
-        joint_axis: The optional 3D axis of rotation or translation of the joint.
+        joint_types: The types of the joints.
+        joint_positions: The positions of the joints.
+        joint_axes: The axes of the joints.
 
     Returns:
-        A tuple containing the homogeneous transformation and the motion subspace.
+        The transforms of the joints.
     """
 
     # Prepare the joint position
-    s = jnp.array(joint_position).astype(float)
+    s = jnp.array(joint_positions).astype(float)
 
     def compute_F() -> tuple[jtp.Matrix, jtp.Array]:
-        return jaxlie.SE3.identity(), jnp.zeros(shape=(6, 1))
+        return jaxlie.SE3.identity()
 
     def compute_R() -> tuple[jtp.Matrix, jtp.Array]:
 
         # Get the additional argument specifying the joint axis.
         # This is a metadata required by only some joint types.
-        axis = jnp.array(joint_axis).astype(float).squeeze()
+        axis = jnp.array(joint_axes).astype(float).squeeze()
 
         pre_H_suc = jaxlie.SE3.from_matrix(
             matrix=jnp.eye(4).at[:3, :3].set(Rotation.from_axis_angle(vector=s * axis))
         )
 
-        S = jnp.vstack(jnp.hstack([jnp.zeros(3), axis]))
-
-        return pre_H_suc, S
+        return pre_H_suc
 
     def compute_P() -> tuple[jtp.Matrix, jtp.Array]:
 
         # Get the additional argument specifying the joint axis.
         # This is a metadata required by only some joint types.
-        axis = jnp.array(joint_axis).astype(float).squeeze()
+        axis = jnp.array(joint_axes).astype(float).squeeze()
 
         pre_H_suc = jaxlie.SE3.from_rotation_and_translation(
             rotation=jaxlie.SO3.identity(),
             translation=jnp.array(s * axis),
         )
 
-        S = jnp.vstack(jnp.hstack([axis, jnp.zeros(3)]))
+        return pre_H_suc
 
-        return pre_H_suc, S
-
-    pre_H_suc, S = jax.lax.switch(
-        index=joint_type,
+    return jax.lax.switch(
+        index=joint_types,
         branches=(
             compute_F,  # JointType.Fixed
             compute_R,  # JointType.Revolute
             compute_P,  # JointType.Prismatic
         ),
-    )
-
-    return pre_H_suc.as_matrix(), S
+    ).as_matrix()