Fix implementation of Transform.inv()

brax/v2/base.py

@@ -130,9 +130,10 @@ def to_local(self, t: 'Transform') -> 'Transform':
     rot = math.quat_mul(math.quat_inv(t.rot), self.rot)
     return Transform(pos=pos, rot=rot)
 
-  def inv(self):
+  def inv(self) -> 'Transform':
     """Invert the transform."""
-    return Transform(pos=-1.0 * self.pos, rot=math.quat_inv(self.rot))
+    inv_rot = math.quat_inv(self.rot)
+    return Transform(pos=math.rotate(-self.pos, inv_rot), rot=inv_rot)
 
   @classmethod
   def create(

brax/v2/math_test.py

@@ -16,7 +16,7 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-from brax.v2 import math
+from brax.v2 import base, math
 import jax
 from jax import numpy as jp
 import numpy as np
@@ -81,5 +81,25 @@ def test_orthogonals(self, i):
     self.assertAlmostEqual(np.abs(a.dot(c)), 0, 6)
 
 
+class TransformInvTest(parameterized.TestCase):
+  """Tests Transform.inv()."""
+
+  @parameterized.parameters(range(100))
+  def test_transform_inv(self, i):
+    np.random.seed(i)
+    pos = np.random.randn(3)
+    rot = np.random.randn(4)
+    rot /= np.linalg.norm(rot)
+
+    a = base.Transform(pos, rot)
+    b = a.inv().do(a)
+    c = a.do(a.inv())
+
+    np.testing.assert_array_almost_equal(b.pos, np.zeros(3))
+    np.testing.assert_array_almost_equal(b.rot, np.array([1.0, 0.0, 0.0, 0.0]))
+    np.testing.assert_array_almost_equal(c.pos, np.zeros(3))
+    np.testing.assert_array_almost_equal(c.rot, np.array([1.0, 0.0, 0.0, 0.0]))
+
+
 if __name__ == '__main__':
   absltest.main()