Question about pipeline init function

[onnoeberhard]

Hi, I am working with the Ant environment and the generalized backend and want to use the Jacobians of the Ant's step function for trajectory optimization. The problem is that the state representation, while compatible with JAX's autodiff, is huge, such that the Jacobians are larger than 1000x1000. It seems to me that the "real" underlying system state is only state.pipeline_state.q and state.pipeline_state.qd (29 dimensions). In the generalized pipeline.py, I found the init function, which seems to transform a $q, \dot q$ representation into the large state representation. My question is if

init(sys, state.q, state.qd) == state

if state is encountered during a rollout of the Ant system. I am interested because I could then use the much smaller Jacobians of the modified dynamics

def f(q, qd, u):
    return step(init(sys, q, qd), u)

Thanks!

[btaba]

Hi @onnoeberhard, state has more than just q, qd for 2 reasons: [1] speed up computation in step (mass_mx_inv not fully being recomputed), and [2] for book-keeping within the codebase. You could conceivably fork the code to rerun init on every timestep, but I'd expect this to be slower for every step computation. Hope that answers the question

[onnoeberhard]

I see, thanks a lot!