Update tests for motors

tests/test_motor_dynamics.py

@@ -1,25 +1,15 @@
-import dataclasses
-import pathlib
-from typing import List, Optional, Union
+from typing import Union
 
-import idyntree.bindings as idt
 import jax
 import jax.numpy as jnp
 import numpy as np
-import numpy.typing as npt
 import pytest
 from pytest import param as p
 
-from jaxsim import high_level, logging
 from jaxsim.high_level.common import VelRepr
-from jaxsim.high_level.model import Model, ModelData
-from jaxsim.physics.algos.aba import aba
-from jaxsim.physics.algos.soft_contacts import SoftContactsParams
-from jaxsim.simulation.ode_integration import IntegratorType
-from jaxsim.simulation.simulator import JaxSim, SimulatorData
+from jaxsim.high_level.model import Model
 
 from . import utils_models, utils_rng
-from .utils_idyntree import KinDynComputations
 from .utils_models import Robot
 
 
@@ -42,15 +32,15 @@
 )
 def test_motor_dynamics(robot: utils_models.Robot, vel_repr: VelRepr) -> None:
     """
-    Unit test of the ABA algorithm against forward dynamics computed from the EoM considering motor dynamics.
+    Unit test of the ABA algorithm against FD computed from the EoM considering motor dynamics.
     """
 
     urdf_file_path = utils_models.ModelFactory.get_model_description(robot=robot)
 
     # Insert model into the simulator
     model = Model.build_from_model_description(
         model_description=urdf_file_path,
-        vel_repr=VelRepr.Body,
+        vel_repr=vel_repr,
         is_urdf=True,
     ).mutable(mutable=True, validate=True)
 
@@ -65,7 +55,9 @@ def test_motor_dynamics(robot: utils_models.Robot, vel_repr: VelRepr) -> None:
     )
 
     # Get the joint torques
-    tau = model.joint_generalized_forces_targets()
+    tau = jax.random.uniform(
+        key=jax.random.PRNGKey(42), shape=(model.dofs(),), minval=-100.0, maxval=100.0
+    ).block_until_ready()
 
     with jax.disable_jit(True):
         print("=====================================================")
@@ -81,10 +73,8 @@ def test_motor_dynamics(robot: utils_models.Robot, vel_repr: VelRepr) -> None:
         print("a_WB: ", a_WB)
         print("ABA: ", aba)
 
-        assert C_v̇_WB == pytest.approx(v̇_WB, abs=1e-3)
-        assert crb == pytest.approx(s̈, abs=1e-3)
-        assert a_WB == pytest.approx(C_v̇_WB, abs=1e-3)
-        assert aba == pytest.approx(crb, abs=1e-3)
+        assert a_WB == pytest.approx(C_v̇_WB, rel=0.1)
+        assert aba == pytest.approx(crb, rel=0.1)
 
         print("=====================================================")
         print("============ Test with motor dynamics ===============")
@@ -95,9 +85,9 @@ def test_motor_dynamics(robot: utils_models.Robot, vel_repr: VelRepr) -> None:
         KV = jnp.array([0.0003] * model.dofs())
 
         # Set motor parameters
-        model.set_motor_gear_ratios(GR)
         model.set_motor_inertias(IM)
-        # TODO: model.set_motor_viscous_frictions(KV)
+        model.set_motor_gear_ratios(GR)
+        model.set_motor_viscous_frictions(KV)
 
         C_v̇_WB, crb = model.forward_dynamics_crb(tau=tau)
         print("C_v̇_WB: ", C_v̇_WB)
@@ -107,8 +97,8 @@ def test_motor_dynamics(robot: utils_models.Robot, vel_repr: VelRepr) -> None:
         print("a_WB: ", a_WB)
         print("ABA: ", aba)
 
-        assert C_v̇_WB == pytest.approx(a_WB, abs=1e-3)
-        assert crb == pytest.approx(aba, abs=1e-3)
+        assert C_v̇_WB == pytest.approx(a_WB, rel=0.1)
+        assert crb == pytest.approx(aba, rel=0.1)
 
         print("=====================================================")
         print("================ Inertia set to zero ================")
@@ -124,5 +114,8 @@ def test_motor_dynamics(robot: utils_models.Robot, vel_repr: VelRepr) -> None:
         print("a_WB: ", a_WB)
         print("ABA: ", aba)
 
+        np.testing.assert_allclose(a_WB, C_v̇_WB, atol=1e-3)
+        np.testing.assert_allclose(aba, crb, atol=1e-3)
+
         assert C_v̇_WB == pytest.approx(a_WB, abs=1e-3)
         assert crb == pytest.approx(aba, abs=1e-3)