Using brax to learn robot dynamics

[shamilmamedov]

Hi there! I have got a few general questions about brax, I would appreciate if experienced users of brax can clarify them for me.

• Is brax the right tool for learning the dynamics of a robot consisting of rigid bodies connected by multi-dof elastic joints?
• Is there a way to get the ODE of the system in the form f(theta, x, u) where theta is the vector of parameters?
• Is there a way to specify arbitrary joint force, e.g., use a neural network as a spring force?

[btaba]

Hi @shamilmamedov , below is an attempt to answer some of the questions (hope it's helpful):

• Brax can help learn dynamics of an articulated rigid body robot with multi-dof joints. We don't currently model muscles or tendons (if that's what you mean by elastic)
• The ODE depends on the backend used, we have multiple backends in brax. f as you describe it, would be pipeline.step here or here. For some background on the generalized physics backend, see chapter 4 here, or for spring see here. I'm sure there are better references but those come top of mind
• To specify an arbitrary joint force, you can use an actuator. When you call pipeline.step(sys, state, action) you can plug in the action from your neural net