diff --git a/tests/test_motor_dynamics.py b/tests/test_motor_dynamics.py index b1d4244f0..000cd2366 100644 --- a/tests/test_motor_dynamics.py +++ b/tests/test_motor_dynamics.py @@ -1,129 +1,127 @@ +import dataclasses import pathlib +from typing import List, Optional, 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 +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 - -sdf_path = ( - pathlib.Path.home() - / "git" - / "element_rl-for-codesign" - / "assets" - / "model" - / "Stickbot.urdf" +from .utils_models import Robot +from . import utils_models, utils_rng +from .utils_idyntree import KinDynComputations + + +@pytest.mark.parametrize( + "robot, vel_repr", + [ + p(*[Robot.DoublePendulum, VelRepr.Inertial], id="DoublePendulum-Inertial"), + p(*[Robot.DoublePendulum, VelRepr.Body], id="DoublePendulum-Body"), + p(*[Robot.DoublePendulum, VelRepr.Mixed], id="DoublePendulum-Mixed"), + p(*[Robot.Ur10, VelRepr.Inertial], id="Ur10-Inertial"), + p(*[Robot.Ur10, VelRepr.Body], id="Ur10-Body"), + p(*[Robot.Ur10, VelRepr.Mixed], id="Ur10-Mixed"), + p(*[Robot.AnymalC, VelRepr.Inertial], id="AnymalC-Inertial"), + p(*[Robot.AnymalC, VelRepr.Body], id="AnymalC-Body"), + p(*[Robot.AnymalC, VelRepr.Mixed], id="AnymalC-Mixed"), + p(*[Robot.Cassie, VelRepr.Inertial], id="Cassie-Inertial"), + p(*[Robot.Cassie, VelRepr.Body], id="Cassie-Body"), + p(*[Robot.Cassie, VelRepr.Mixed], id="Cassie-Mixed"), + ], ) +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. + """ -simulator = JaxSim.build( - step_size=1e-3, - steps_per_run=1, - velocity_representation=high_level.model.VelRepr.Body, - integrator_type=IntegratorType.EulerSemiImplicit, - simulator_data=SimulatorData( - contact_parameters=SoftContactsParams(K=1e6, D=1.5e4, mu=0.5), - ), -).mutable(validate=False) - -# Insert model into the simulator -model = simulator.insert_model_from_description(model_description=sdf_path).mutable( - validate=True -) -model.reduce( - considered_joints=[ - "r_shoulder_pitch", - "r_shoulder_roll", - "r_shoulder_yaw", - "r_elbow", - "l_shoulder_pitch", - "l_shoulder_roll", - "l_shoulder_yaw", - "l_elbow", - "r_hip_pitch", - "r_hip_roll", - "r_hip_yaw", - "r_knee", - "r_ankle_pitch", - "r_ankle_roll", - "l_hip_pitch", - "l_hip_roll", - "l_hip_yaw", - "l_knee", - "l_ankle_pitch", - "l_ankle_roll", - "torso_roll", - "torso_yaw", - ] -) + 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, + is_urdf=True, + ).mutable(mutable=True, validate=True) + + # Initialize the model with a random state + model.data.model_state = utils_rng.random_physics_model_state( + physics_model=model.physics_model + ) + + # Initialize the model with a random input + model.data.model_input = utils_rng.random_physics_model_input( + physics_model=model.physics_model + ) + + # Get the joint torques + tau = model.joint_generalized_forces_targets() -print("Model links: ", model.link_names()) -tau = jnp.array([1.0] * model.dofs()) + with jax.disable_jit(True): + print("=====================================================") + print("============ Test without motor dynamics ============") + print("=====================================================") -with jax.disable_jit(): - print("=====================================================") - print("============ Test without motor dynamics ============") - print("=====================================================") + C_v̇_WB, crb = model.forward_dynamics_crb(tau=tau) - C_v̇_WB, crb = model.forward_dynamics_crb(tau=tau) + print("C_v̇_WB: ", C_v̇_WB) + print("CRB: ", crb) - print("C_v̇_WB: ", C_v̇_WB) - print("CRB: ", crb) + a_WB, aba = model.forward_dynamics_aba(tau=tau) + print("a_WB: ", a_WB) + print("ABA: ", aba) - a_WB, aba = model.forward_dynamics_aba(tau=tau) - 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) - try: - np.testing.assert_allclose(crb, aba, rtol=0.5) - np.testing.assert_allclose(C_v̇_WB, a_WB, rtol=0.5) - except Exception as e: - print(e) + print("=====================================================") + print("============ Test with motor dynamics ===============") + print("=====================================================") - print("=====================================================") - print("============ Test with motor dynamics ===============") - print("=====================================================") + IM = jnp.array([5e-4] * model.dofs()) + GR = jnp.array([100.0] * model.dofs()) + KV = jnp.array([0.0003] * model.dofs()) - # Set motor parameters - model.set_motor_gear_ratios(jnp.array([1 / 100.0] * model.dofs())) - model.set_motor_inertias(jnp.array([5e-4] * model.dofs())) - model.set_motor_viscous_frictions(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) - C_v̇_WB, crb = model.forward_dynamics_crb(tau=tau) - print("C_v̇_WB: ", C_v̇_WB) - print("CRB: ", crb) + C_v̇_WB, crb = model.forward_dynamics_crb(tau=tau) + print("C_v̇_WB: ", C_v̇_WB) + print("CRB: ", crb) - a_WB, aba = model.forward_dynamics_aba(tau=tau) - print("a_WB: ", a_WB) - print("ABA: ", aba) + a_WB, aba = model.forward_dynamics_aba(tau=tau) + print("a_WB: ", a_WB) + print("ABA: ", aba) - try: - np.testing.assert_allclose(crb, aba, rtol=0.5) - np.testing.assert_allclose(C_v̇_WB, a_WB, rtol=0.5) - except Exception as e: - print(e) + assert C_v̇_WB == pytest.approx(a_WB, abs=1e-3) + assert crb == pytest.approx(aba, abs=1e-3) - print("=====================================================") - print("================ Inertia set to zero ================") - print("=====================================================") + print("=====================================================") + print("================ Inertia set to zero ================") + print("=====================================================") - # Set motor inertia to zero and perform FD with CRB and ABA - model.set_motor_inertias(jnp.array([0.0] * model.dofs())) + model.set_motor_inertias(jnp.array([0.0] * model.dofs())) - C_v̇_WB, crb = model.forward_dynamics_crb(tau=tau) - print("C_v̇_WB: ", C_v̇_WB) - print("CRB: ", crb) + C_v̇_WB, crb = model.forward_dynamics_crb(tau=tau) + print("C_v̇_WB: ", C_v̇_WB) + print("CRB: ", crb) - a_WB, aba = model.forward_dynamics_aba(tau=tau) - print("a_WB: ", a_WB) - print("ABA: ", aba) + a_WB, aba = model.forward_dynamics_aba(tau=tau) + print("a_WB: ", a_WB) + print("ABA: ", aba) - try: - np.testing.assert_allclose(crb, aba, rtol=0.5) - np.testing.assert_allclose(C_v̇_WB, a_WB, rtol=0.5) - except Exception as e: - print(e) + assert C_v̇_WB == pytest.approx(a_WB, abs=1e-3) + assert crb == pytest.approx(aba, abs=1e-3)