Cache base_velocity, base_orientation and base_transform

ami-iit · Jan 15, 2025 · dc458fa · dc458fa
1 parent 00c7b84
commit dc458fa
Show file tree

Hide file tree

Showing 19 changed files with 274 additions and 261 deletions.
diff --git a/examples/jaxsim_for_robot_controllers.ipynb b/examples/jaxsim_for_robot_controllers.ipynb
@@ -256,7 +256,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",
@@ -345,8 +345,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",
-    "    ṡ = data.joint_velocities()\n",
+    "    s = data.joint_positions\n",
+    "    ṡ = data.joint_velocities\n",
     "\n",
     "    # Compute the actuated joint torques.\n",
     "    s_star = -kp * (s - s_des) - kd * (ṡ - s_dot_des)\n",
@@ -405,7 +405,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",

diff --git a/src/jaxsim/api/com.py b/src/jaxsim/api/com.py
@@ -27,7 +27,7 @@ def com_position(
     m = js.model.total_mass(model=model)
 
     W_H_L = data.kyn_dyn.forward_kinematics
-    W_H_B = data.base_transform()
+    W_H_B = data.kyn_dyn.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.kyn_dyn.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.kyn_dyn.base_transform
     W_p_CoM = com_position(model=model, data=data)
 
     match data.velocity_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.kyn_dyn.base_transform.at[0:3].set(W_p_CoM)
             ).T
 
             GB_ḣ_bias = GB_Xf_W @ W_ḣ_bias

diff --git a/src/jaxsim/api/contact.py b/src/jaxsim/api/contact.py
@@ -636,20 +636,20 @@ def compute_Ṫ(model: js.model.JaxSimModel, Ẋ: jtp.Matrix) -> jtp.Matrix:
             Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_W)
 
         case VelRepr.Body:
-            W_H_B = data.base_transform()
+            W_H_B = data.kyn_dyn.base_transform
             W_X_B = Adjoint.from_transform(transform=W_H_B)
-            B_v_WB = data.base_velocity()
+            B_v_WB = data.kyn_dyn.base_velocity
             B_vx_WB = Cross.vx(B_v_WB)
             W_Ẋ_B = W_X_B @ B_vx_WB
 
             T = compute_T(model=model, X=W_X_B)
             Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_B)
 
         case VelRepr.Mixed:
-            W_H_B = data.base_transform()
+            W_H_B = data.kyn_dyn.base_transform
             W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))
             W_X_BW = Adjoint.from_transform(transform=W_H_BW)
-            BW_v_WB = data.base_velocity()
+            BW_v_WB = data.kyn_dyn.base_velocity
             BW_v_W_BW = BW_v_WB.at[3:6].set(jnp.zeros(3))
             BW_vx_W_BW = Cross.vx(BW_v_W_BW)
             W_Ẋ_BW = W_X_BW @ BW_vx_W_BW