Save some kindyn computation in JaxSimModelData

examples/jaxsim_as_physics_engine.ipynb

@@ -299,7 +299,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 +398,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 +412,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

@@ -239,6 +239,8 @@
     "# Reset the state to the random joint positions.\n",
     "data = data_zero.reset_joint_positions(positions=random_joint_positions)\n",
     "\n",
+    "# Update the kyn_dyn cache.\n",
+    "data = data.update_kyn_dyn(model=model)\n",
     "\n",
     "for _ in T:\n",
     "\n",

src/jaxsim/api/com.py

@@ -26,7 +26,7 @@ def com_position(
 
     m = js.model.total_mass(model=model)
 
-    W_H_L = js.model.forward_kinematics(model=model, data=data)
+    W_H_L = data.kyn_dyn.forward_kinematics
     W_H_B = data.base_transform()
     B_H_W = jaxsim.math.Transform.inverse(transform=W_H_B)
 
@@ -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.kyn_dyn.forward_kinematics
 
     # Compute the bias acceleration of all links by zeroing the generalized velocity
     # in the active representation.

src/jaxsim/api/common.py

@@ -232,3 +232,33 @@ def other_representation_to_inertial(
 
             case _:
                 raise ValueError(other_representation)
+
+
+def convert_mass_matrix(
+    M: jtp.Matrix,
+    base_transform: jtp.Matrix,
+    dofs: jtp.Int,
+    velocity_representation: VelRepr,
+):
+    # TODO (flferretti): perform the velocity representation conversion here.
+    return M
+
+
+def convert_jacobian(
+    J: jtp.Matrix,
+    base_transform: jtp.Matrix,
+    dofs: jtp.Int,
+    velocity_representation: VelRepr,
+):
+    # TODO (flferretti): save actual Jacobian instead of full doubly left and perform conversion.
+    return J
+
+
+def convert_jacobian_derivative(
+    Jd: jtp.Matrix,
+    base_transform: jtp.Matrix,
+    dofs: jtp.Int,
+    velocity_representation: VelRepr,
+):
+    # TODO (flferretti): save actual Jacobian derivative instead of full doubly left and perform conversion.
+    return Jd

src/jaxsim/api/contact.py

@@ -40,15 +40,8 @@ def collidable_point_kinematics(
     # Switch to inertial-fixed since the RBDAs expect velocities in this representation.
     with data.switch_velocity_representation(VelRepr.Inertial):
 
-        W_p_Ci, W_ṗ_Ci = jaxsim.rbda.collidable_points.collidable_points_pos_vel(
-            model=model,
-            base_position=data.base_position(),
-            base_quaternion=data.base_orientation(dcm=False),
-            joint_positions=data.joint_positions(model=model),
-            base_linear_velocity=data.base_velocity()[0:3],
-            base_angular_velocity=data.base_velocity()[3:6],
-            joint_velocities=data.joint_velocities(model=model),
-        )
+        W_p_Ci = data.kyn_dyn.collidable_point_positions
+        W_ṗ_Ci = data.kyn_dyn.collidable_point_velocities
 
     return W_p_Ci, W_ṗ_Ci
 
@@ -460,7 +453,7 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
     )[indices_of_enabled_collidable_points]
 
     # Get the transforms of the parent link of all collidable points.
-    W_H_L = js.model.forward_kinematics(model=model, data=data)[
+    W_H_L = data.kyn_dyn.forward_kinematics[
         parent_link_idx_of_enabled_collidable_points
     ]
 
@@ -612,7 +605,7 @@ def jacobian_derivative(
     ]
 
     # Get the transforms of all the parent links.
-    W_H_Li = js.model.forward_kinematics(model=model, data=data)
+    W_H_Li = data.kyn_dyn.forward_kinematics
 
     # =====================================================
     # Compute quantities to adjust the input representation