Cache base_velocity, base_orientation and base_transform

ami-iit · Jan 15, 2025 · 0e79852 · 0e79852
1 parent 295052b
commit 0e79852
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 35 deletions.
diff --git a/src/jaxsim/api/model.py b/src/jaxsim/api/model.py
@@ -704,6 +704,7 @@ def generalized_free_floating_jacobian_derivative(
     # by zeroing the columns not in the path π_B(L) using the boolean κ(i).
     κb = model.kin_dyn_parameters.support_body_array_bool
 
+    # Compute the base transform.
     W_H_B = data.kyn_dyn.base_transform
 
     # We add the 5 columns of ones to the Jacobian derivative to account for the

diff --git a/tests/test_api_model.py b/tests/test_api_model.py
@@ -283,49 +283,50 @@ def test_model_rbda(
 
     _, subkey = jax.random.split(prng_key, num=2)
 
-    data = js.data.random_model_data(
-        model=model, key=subkey, velocity_representation=velocity_representation
-    )
+    with jax.disable_jit():
+        data = js.data.random_model_data(
+            model=model, key=subkey, velocity_representation=velocity_representation
+        )
 
-    kin_dyn = utils_idyntree.build_kindyncomputations_from_jaxsim_model(
-        model=model, data=data
-    )
+        kin_dyn = utils_idyntree.build_kindyncomputations_from_jaxsim_model(
+            model=model, data=data
+        )
 
-    # =====
-    # Tests
-    # =====
+        # =====
+        # Tests
+        # =====
 
-    # Support both fixed-base and floating-base models by slicing the first six rows.
-    sl = np.s_[0:] if model.floating_base() else np.s_[6:]
+        # Support both fixed-base and floating-base models by slicing the first six rows.
+        sl = np.s_[0:] if model.floating_base() else np.s_[6:]
 
-    # Mass matrix
-    M_idt = kin_dyn.mass_matrix()
-    M_js = js.model.free_floating_mass_matrix(model=model, data=data)
-    assert pytest.approx(M_idt[sl, sl]) == M_js[sl, sl]
+        # Mass matrix
+        M_idt = kin_dyn.mass_matrix()
+        M_js = js.model.free_floating_mass_matrix(model=model, data=data)
+        assert pytest.approx(M_idt[sl, sl]) == M_js[sl, sl]
 
-    # Gravity forces
-    g_idt = kin_dyn.gravity_forces()
-    g_js = js.model.free_floating_gravity_forces(model=model, data=data)
-    assert pytest.approx(g_idt[sl]) == g_js[sl]
+        # Gravity forces
+        g_idt = kin_dyn.gravity_forces()
+        g_js = js.model.free_floating_gravity_forces(model=model, data=data)
+        assert pytest.approx(g_idt[sl]) == g_js[sl]
 
-    # Bias forces
-    h_idt = kin_dyn.bias_forces()
-    h_js = js.model.free_floating_bias_forces(model=model, data=data)
-    assert pytest.approx(h_idt[sl]) == h_js[sl]
+        # Bias forces
+        h_idt = kin_dyn.bias_forces()
+        h_js = js.model.free_floating_bias_forces(model=model, data=data)
+        assert pytest.approx(h_idt[sl]) == h_js[sl]
 
-    # Forward kinematics
-    HH_js = js.model.forward_kinematics(model=model, data=data)
-    HH_idt = jnp.stack(
-        [kin_dyn.frame_transform(frame_name=name) for name in model.link_names()]
-    )
-    assert pytest.approx(HH_idt) == HH_js
+        # Forward kinematics
+        HH_js = js.model.forward_kinematics(model=model, data=data)
+        HH_idt = jnp.stack(
+            [kin_dyn.frame_transform(frame_name=name) for name in model.link_names()]
+        )
+        assert pytest.approx(HH_idt) == HH_js
 
-    # Bias accelerations
-    Jν_js = js.model.link_bias_accelerations(model=model, data=data)
-    Jν_idt = jnp.stack(
-        [kin_dyn.frame_bias_acc(frame_name=name) for name in model.link_names()]
-    )
-    assert pytest.approx(Jν_idt) == Jν_js
+        # Bias accelerations
+        Jν_js = js.model.link_bias_accelerations(model=model, data=data)
+        Jν_idt = jnp.stack(
+            [kin_dyn.frame_bias_acc(frame_name=name) for name in model.link_names()]
+        )
+        assert pytest.approx(Jν_idt) == Jν_js
 
 
 def test_model_jacobian(