ami-iit / jaxsim

A differentiable physics engine and multibody dynamics library for control and robot learning.
https://jaxsim.readthedocs.io/
BSD 3-Clause "New" or "Revised" License
57 stars 9 forks source link

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

Closed diegoferigo closed 2 weeks ago

diegoferigo commented 2 weeks ago

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:

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/