Always forward to the integrator the model and data passed to jaxsim.api.model.step

[diegoferigo]

This modification is necessary to re-use the same jit-compiled jaxsim.api.model.step function and the same integrator (closed over an initial model) on a modified model having the same structure of the initial model.

It enables running simulations with parametric models having constant pytree structure without incurring in jit recompilations.

For more details, refer to the following logic:

• jaxsim.api.model.step.
• jaxsim.integrators.common.Integrator.step calling jaxsim.integrators.common.Integrator.__call__. For example, our fixed-step integrators pass the kwargs to the system dynamics here.
• The system dynamic function considered by the integrator is created in jaxsim.api.ode.wrap_system_dynamics_for_integration. By default the integrator operating on this dynamics considers the model and data (mainly, the parameters stored in data that are not data.state) that are passed during the initialization of the integrator. This PR always overrides model and data with those passed to jaxsim.api.model.step.

Note: furthermore, if, for any reason, the user wants to pass to the integrator a different model or data, they can do it as follows:

_ = js.model.step(
    model=model,
    data=data,
    dt=dt,
    integrator=integrator,
    integrator_state=integrator_state,
    # Override model and data:
    integrator_kwargs=dict(model=other_model, data=other_data),
)

Related to #101 and #120.

---

📚 Documentation preview 📚: https://jaxsim--183.org.readthedocs.build//183/