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.
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),
)
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
callingjaxsim.integrators.common.Integrator.__call__
. For example, our fixed-step integrators pass thekwargs
to the system dynamics here.jaxsim.api.ode.wrap_system_dynamics_for_integration
. By default the integrator operating on this dynamics considers themodel
anddata
(mainly, the parameters stored in data that are notdata.state
) that are passed during the initialization of the integrator. This PR always overridesmodel
anddata
with those passed tojaxsim.api.model.step
.Note: furthermore, if, for any reason, the user wants to pass to the integrator a different
model
ordata
, they can do it as follows:Related to #101 and #120.
📚 Documentation preview 📚: https://jaxsim--183.org.readthedocs.build//183/