Fix implementation of Transform.inv()

google / brax

Massively parallel rigidbody physics simulation on accelerator hardware.

Apache License 2.0

2.25k stars 249 forks source link

Fix implementation of Transform.inv() #316

Closed allanzhao closed 22 hours ago

allanzhao commented 1 year ago

I've been porting some of my research code to Brax v2, and the performance is very good so far! Our code uses rigid body transforms extensively, and we noticed that Transform.inv() seems to differ from the textbook formula (as mentioned in #297).

This pull request: 1) changes the implementation of brax.v2.base.Transform.inv() to match tiny-differentiable-simulator and MuJoCo 2) adds a test case at the end of brax/v2/math_test.py (please feel free to move to an appropriate place)

btaba commented 1 year ago

Hi @allanzhao ! Thanks for the PR and documenting the bug! We removed Transform.inv() in favor of inv_do https://github.com/google/brax/blob/9d5a88ef69606966ad075a7b043b5d1541342991/brax/v2/base.py#L127-L129

the reason being that Transform.inv(), as implemented here, is wrt other Transforms, not necessarily wrt a Motion or a Force which is a bit confusing. We were only using Transform.inv_do(Motion), but if you want to change this into a Transform.inv_do(Transform) operator then that would be great!