[Sprint] Avoid evaluation of branches leading to NaNs in the optimizer

docs/guide/configuration.rst

@@ -61,7 +61,7 @@ The logging and exceptions configurations is controlled by the following environ
 
   *Default:* ``DEBUG`` for development, ``WARNING`` for production.
 
-- ``JAXSIM_DISABLE_EXCEPTIONS``: Disables the runtime checks and exceptions.
+- ``JAXSIM_ENABLE_EXCEPTIONS``: Enables the runtime checks and exceptions. Note that enabling exceptions might lead to device-to-host transfer of data, increasing the computational time required.
 
   *Default:* ``False``.
 

environment.yml

@@ -27,6 +27,7 @@ dependencies:
   # [testing]
   - idyntree >= 12.2.1
   - pytest
+  - pytest-benchmark
   - pytest-icdiff
   - robot_descriptions
   - icub-models

examples/jaxsim_as_multibody_dynamics_library.ipynb

@@ -368,7 +368,7 @@
     "# Alternatively, the same transform can be extracted from the tensor of\n",
     "# all link transforms.\n",
     "assert jnp.allclose(\n",
-    "    W_H_L, js.model.forward_kinematics(model=model, data=data)[link_index]\n",
+    "    W_H_L, data.link_transforms[link_index]\n",
     ")\n",
     "\n",
     "print(f\"Transform of '{link_name}': shape={W_H_L.shape}\\n{W_H_L}\")"
@@ -782,12 +782,11 @@
     ")\n",
     "\n",
     "# Compute the 3D gravity vector and the total mass of the robot.\n",
-    "W_g = data.gravity\n",
     "m = js.model.total_mass(model=model)\n",
     "\n",
     "# The centroidal dynamics can be computed as follows.\n",
     "G_ḣ = 0\n",
-    "G_ḣ += m * jnp.hstack([W_g, jnp.zeros(3)])\n",
+    "G_ḣ += m * jnp.hstack([0, 0, model.gravity, 0, 0, 0])\n",
     "G_ḣ += jnp.einsum(\"c66,c6->6\", G_Xf_C, jnp.hstack([C_fl, jnp.zeros_like(C_fl)]))\n",
     "print(f\"G_ḣ:    shape={G_ḣ.shape}\")"
    ]

examples/jaxsim_as_physics_engine.ipynb

@@ -175,7 +175,7 @@
     "\n",
     "# Simulate\n",
     "for _t in T:\n",
-    "    data, _ = js.model.step(\n",
+    "    data = js.model.step(\n",
     "        model=model,\n",
     "        data=data,\n",
     "        link_forces=None,\n",
@@ -239,7 +239,7 @@
     "\n",
     "data = data_batch_t0\n",
     "for _t in T:\n",
-    "    data, _ = step_parallel(model, data)"
+    "    data = step_parallel(model, data)"
    ]
   }
  ],
@@ -266,7 +266,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.0"
+   "version": "3.13.1"
   }
  },
  "nbformat": 4,

examples/jaxsim_as_physics_engine_advanced.ipynb

@@ -63,7 +63,6 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import jaxsim.api as js\n",
-    "import jaxsim\n",
     "import rod\n",
     "from jaxsim import logging\n",
     "from rod.builder.primitives import SphereBuilder\n",
@@ -163,7 +162,6 @@
     "model = js.model.JaxSimModel.build_from_model_description(\n",
     "    model_description=model_sdf_string,\n",
     "    time_step=0.001,\n",
-    "    integrator=jaxsim.integrators.fixed_step.Heun2,\n",
     ")\n",
     "\n",
     "# Create the data of a single model.\n",
@@ -299,7 +297,7 @@
     "    )\n",
     ")(jnp.vstack(subkeys))\n",
     "\n",
-    "print(\"W_p_B(t0)=\\n\", data_batch_t0.base_position()[0:10])"
+    "print(\"W_p_B(t0)=\\n\", data_batch_t0.base_position[0:10])"
    ]
   },
   {
@@ -398,7 +396,7 @@
     "# This operation is called 'tree transpose' in JAX.\n",
     "data_trajectory = jax.tree.map(lambda *leafs: jnp.stack(leafs), *data_trajectory_list)\n",
     "\n",
-    "print(f\"W_p_B: shape={data_trajectory.base_position().shape}\")"
+    "print(f\"W_p_B: shape={data_trajectory.base_position.shape}\")"
    ]
   },
   {
@@ -412,7 +410,7 @@
     "import matplotlib.pyplot as plt\n",
     "\n",
     "\n",
-    "plt.plot(T, data_trajectory.base_position()[:, 0:5, 2])\n",
+    "plt.plot(T, data_trajectory.base_position[:, 0:5, 2])\n",
     "plt.grid(True)\n",
     "plt.xlabel(\"Time [s]\")\n",
     "plt.ylabel(\"Height [m]\")\n",

examples/jaxsim_for_robot_controllers.ipynb

@@ -254,7 +254,7 @@
     "\n",
     "    # Update the MuJoCo data.\n",
     "    mj_model_helper.set_joint_positions(\n",
-    "        positions=data.joint_positions(), joint_names=model.joint_names()\n",
+    "        positions=data.joint_positions, joint_names=model.joint_names()\n",
     "    )\n",
     "\n",
     "    # Record a new video frame.\n",
@@ -343,8 +343,8 @@
     "    Mss = js.model.free_floating_mass_matrix(model=model, data=data)[6:, 6:]\n",
     "\n",
     "    # Get the current joint positions and velocities.\n",
-    "    s = data.joint_positions()\n",
-    "    ṡ = data.joint_velocities()\n",
+    "    s = data.joint_positions\n",
+    "    ṡ = data.joint_velocities\n",
     "\n",
     "    # Compute the actuated joint torques.\n",
     "    s_star = -kp * (s - s_des) - kd * (ṡ - s_dot_des)\n",
@@ -403,7 +403,7 @@
     "\n",
     "    # Update the MuJoCo data.\n",
     "    mj_model_helper.set_joint_positions(\n",
-    "        positions=data.joint_positions(), joint_names=model.joint_names()\n",
+    "        positions=data.joint_positions, joint_names=model.joint_names()\n",
     "    )\n",
     "\n",
     "    # Record a new video frame.\n",

src/jaxsim/__init__.py

@@ -114,5 +114,5 @@ def _get_default_logging_level(env_var: str) -> logging.LoggingLevel:
 del _is_editable
 
 from . import terrain  # isort:skip
-from . import api, integrators, logging, math, rbda
+from . import api, logging, math, rbda
 from .api.common import VelRepr

src/jaxsim/api/__init__.py

@@ -1,13 +1,15 @@
 from . import common  # isort:skip
 from . import model, data  # isort:skip
 from . import (
+    actuation_model,
     com,
     contact,
+    contact_model,
     frame,
+    integrators,
     joint,
     kin_dyn_parameters,
     link,
     ode,
-    ode_data,
     references,
 )

src/jaxsim/api/actuation_model.py

@@ -0,0 +1,96 @@
+import jax.numpy as jnp
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+
+
+def compute_resultant_torques(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    joint_force_references: jtp.Vector | None = None,
+) -> jtp.Vector:
+    """
+    Compute the resultant torques acting on the joints.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        joint_force_references: The joint force references to apply.
+
+    Returns:
+        The resultant torques acting on the joints.
+    """
+
+    # 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.joint_positions
+            - model.kin_dyn_parameters.joint_parameters.position_limits_min,
+            max=0.0,
+        )
+
+        upper_violation = jnp.clip(
+            data.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.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.joint_velocities)
+            + jnp.diag(kv) @ data.joint_velocities
+        )
+
+    # ===============================
+    # Compute the total joint forces.
+    # ===============================
+
+    τ_total = τ_references + τ_friction + τ_position_limit
+
+    return τ_total

src/jaxsim/api/com.py

@@ -26,8 +26,8 @@ def com_position(
 
     m = js.model.total_mass(model=model)
 
-    W_H_L = js.model.forward_kinematics(model=model, data=data)
-    W_H_B = data.base_transform()
+    W_H_L = data.link_transforms
+    W_H_B = data.base_transform
     B_H_W = jaxsim.math.Transform.inverse(transform=W_H_B)
 
     def B_p̃_LCoM(i) -> jtp.Vector:
@@ -134,7 +134,7 @@ def centroidal_momentum_jacobian(
         model=model, data=data, output_vel_repr=VelRepr.Body
     )
 
-    W_H_B = data.base_transform()
+    W_H_B = data.base_transform
     B_H_W = jaxsim.math.Transform.inverse(W_H_B)
 
     W_p_CoM = com_position(model=model, data=data)
@@ -172,7 +172,7 @@ def locked_centroidal_spatial_inertia(
     with data.switch_velocity_representation(VelRepr.Body):
         B_Mbb_B = js.model.locked_spatial_inertia(model=model, data=data)
 
-    W_H_B = data.base_transform()
+    W_H_B = data.base_transform
     W_p_CoM = com_position(model=model, data=data)
 
     match data.velocity_representation:
@@ -269,7 +269,7 @@ def bias_acceleration(
     """
 
     # Compute the pose of all links with forward kinematics.
-    W_H_L = js.model.forward_kinematics(model=model, data=data)
+    W_H_L = data.link_transforms
 
     # Compute the bias acceleration of all links by zeroing the generalized velocity
     # in the active representation.
@@ -411,7 +411,7 @@ def bias_momentum_derivative_term(
         case VelRepr.Body:
 
             GB_Xf_W = jaxsim.math.Adjoint.from_transform(
-                transform=data.base_transform().at[0:3].set(W_p_CoM)
+                transform=data.base_transform.at[0:3].set(W_p_CoM)
             ).T
 
             GB_ḣ_bias = GB_Xf_W @ W_ḣ_bias