Brax contact modeling results in non-physical behavior (and non-physical gradients)

[dawsonc]

Hi,

I'm noticing an issue in how Brax handles contact in both the pbd and legacy_spring dynamics modes, where both the auto-diff computed gradient and the numerically-estimated gradient will not agree with the physical behavior of the system.

The setup is in this Colab notebook. I created a simple environment with an elastic ball, a ground plane, and no gravity. The ball is initialized with a constant downwards velocity and the simulation is run for a fixed length (1 second; set to be long enough that the ball bounces once). I then vary the initial height of the ball between 2.9 and 3 meters and observe the final height of the ball.

Thinking about what should happen physically: with no gravity, if we raise the initial height of the ball, then it should take longer to reach the floor and bounce, thus spending less time in the "upwards moving" phase and reaching a lower final height. However, although this trend (higher initial height -> lower final height) holds on average, the final height is not continuous when we simulate it using Brax!

Here's the plot of final height vs initial height with the pbd backend:

And here's a plot of the final height vs initial height with the legacy_spring backend:

If we were to try to compute the gradient of the final height w.r.t. the initial height (either with jax.grad or with some numerical difference scheme), then the gradient would be zero almost everywhere (for pbd) or opposite the long-run trend (for legacy_spring)!

This seems to be the same issue as reported in this paper (section 4.3). Here's the key figure (note how it matches Brax's behavior for the spring case)

There's also a video showing this issue here.

That paper suggests adding time-of-collision detection as a remedy; is that something that could be added to Brax? This is in the critical path of my research, and so I would be willing to help implement this feature, but I might need some pointers for where to start.

Thanks for your help! I like Brax and would love to help out.

P.S. it seems like many people have struggled to get analytical policy gradient-based algorithms to work. Could this be the reason?

[cdfreeman-google]

Hello! Thank you for the extensive analysis!

Yes, time of impact has been on our radar as something that would be a nice-to-have for a while (see this old issue: https://github.com/google/brax/issues/73). While we haven't gotten around to it, I don't think it would be a huge undertaking, and I can definitely provide some pointers if you'd like to pursue it in a PR.

For a sketch of how to go about this, you would want to modify the collision functions to return the time of impact delta (call it toi), and then pass this along to the integrator to partially step time. The slickest implementation I'm aware of that would be a good code reference is in the difftaichi repo here: https://github.com/taichi-dev/difftaichi/blob/master/examples/rigid_body.py (look for toi references in there).

This is definitely one of the big problems plaguing most simulators that folks try to take gradients through. There's an additional problem that the underlying dynamics are just fundamentally chaotic a lot of the time, which you ultimately won't be able to avoid, but this might be able to extend the "longest trajectory that it is possible to optimize" out a little ways.

[dawsonc]

Alright; thanks for the pointers! I'll take a look over the next few weeks and see how it goes.

[EelcoHoogendoorn]

The simplest solution is to have soft contact, which resolve over multiple timesteps. If the physics of your problem actually do require hard contacts you are in a bit of trouble of course, unless you are very patient in substepping your physics. But for some problems soft contacts are actually appropriate or acceptable; and the least-invasive change would be to just accomodate that with a parameter.

[dawsonc]

Sorry -- this completely fell of my radar (my research ended up moving in a different direction that did not involve Brax). I'm likely not going to have the bandwidth to fix this in the near term; if anyone else wants to give it a go, feel free.

@EelcoHoogendoorn true enough that integrating the collision over a couple of timesteps would help resolve this (but may require some problem-specific tuning to make sure you're actually getting enough timesteps during the contact).

[diegoferigo]

The simplest solution is to have soft contact, which resolve over multiple timesteps. If the physics of your problem actually do require hard contacts you are in a bit of trouble of course, unless you are very patient in substepping your physics. But for some problems soft contacts are actually appropriate or acceptable; and the least-invasive change would be to just accomodate that with a parameter.

I want to chime in for a quick comment on soft contacts. I am the author of ami-iit/jaxsim, a physics engine similar to Brax (actually, originally inspired by it) that got extended as multibody dynamics library.

From the very beginning, JaxSim simulates free-floating robots using a soft contacts model together with a simple point-surface collision detection (useful for locomotion application). The project is still under development, there are not yet many examples although you can see the contact model running in https://github.com/ami-iit/jaxsim/pull/125.

We are still in the process of developing and running closed-loop architectures developed in JAX and simulated in JaxSim, therefore I can only provide initial feedback on that if needed - probably out of scope of this issue. Here below, however, I leave some comment on our experience with soft contacts so far:

• Our soft contacts model extends the state of robot dynamics (the ODE that is integrated numerically) with the material deformation caused by contacts. With this additional state, the soft contacts model can be treated as an instantaneous function of the state (therefore the contact forces can be computed analytically).
• Being an instantaneous function of the state, there is no need to design nor run any iterative optimization routine to compute the contact forces.
• Gradients seem to work correctly, however we have never yet used JaxSim for any gradient-based optimization.
• The resulting ODE, depending on the soft-contacts parameters, can become particularly stiff. Fixed-step integrators may need a tiny time-step to obtain a stable simulation - that can become particularly slow. Things can be mitigated with variable-step integrators, paying the price of losing reverse-mode AD - at least in their basic implementation.
• Last but not least, the soft-contacts parameters are not easy to tune. We provide a good starting point by estimating them depending on the model weight and supporting contact points, though manual adjustment is often needed.

Finally, keep also in mind that the non-zero penetration that all soft contacts models introduce may challenge traditional reactive/predictive whole-body controllers, that typically model holonomic constraints of contact points assuming rigid contacts.

I hope all of this could be useful in case anyone wants to explore the introduction of soft contacts in Brax.

[EelcoHoogendoorn]

@diegoferigo nice to see more people doing simulation in Jax. 6d inertia matrices look nicer still when viewed as inertia maps of bivectors tho :)