Potentially naive question: are collision callbacks a thing/possible in brax?

[guydav]

Hello,

Apologies in advance for the potentially naive question, but googling and some skimming of the source code didn't seem to provide a clear answer:

In other physics engines I've used before (e.g., matter.js or Pymunk), it's trivial to register callbacks for collision events between specific objects to modify their behavior. The use cases usually modify collision behavior, e.g., by applying an additional force on one of the colliding objects.

I don't immediately see a way to do this in brax -- is this supported by any of the pipelines implemented in brax? If not, how outrageously out of the trodden path of brax would it be to implement something of this nature? I'm asking because I might have to swap out physics engines for something I'm working on. I heard about brax, and wanted to try to use it, but I have a bunch of behaviors I've found easiest to conceptualize as custom collision handlers.

Thank you!

[erikfrey]

We don't offer explicit hooks, but we have designed the API to be easily decomposable to inject your own physics if you wish.

Consider your use case, adding an external force. You could take MJX's forward function and add some extra forces after detecting contacts:

def my_cool_forward(m: Model, d: Data) -> Data:
  """My cool forward dynamics that adds extra forces to contact."""
  d = fwd_position(m, d)
  # add forces based on contacts
  d = d.replace(xfrc_applied=sum_contacts_to_body(d, force=0.1))
  d = fwd_velocity(m, d)
  d = fwd_actuation(m, d)
  d = fwd_acceleration(m, d)

  if d.efc_J.size == 0:
    d = d.replace(qacc=d.qacc_smooth)
    return d

  d = named_scope(solver.solve)(m, d)

  return d

def my_cool_step(m: Model, d: Data) -> Data:
  """Advance simulation."""
  d = my_cool_forward(m, d)
  d = euler(m, d)

  return d    

Then you could pipe that into a custom brax.PipelineEnv and off you go. It's not the cleanest but it would work.

[guydav]

Thanks, @erikfrey! This is helpful and I'll follow up if and when I get to give it a go. One more question, if I may -- is there a standard way to constrain movement to two dimensions? Would it be something similar to modify the forward pass to eliminate any forces/acceleration/velocity in the Z (or any other) direction?

[erikfrey]

Sure, see the target defined in the reacher environment for an example of a body that is constrained to movement in two directions.

[guydav]

Got it, thanks again!