princeton-vl / lietorch

BSD 3-Clause "New" or "Revised" License
661 stars 46 forks source link

Calling .quaternion() produces runtime error #10

Open TobiasJacob opened 2 years ago

TobiasJacob commented 2 years ago

This is a minimal example demonstrating the issue:

from lietorch import SE3
ex = SE3.Identity(1)
ex.quaternion()

results in

NameError: name 'Quat' is not defined
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
/tmp/ipykernel_88211/189352308.py in <module>
      1 from lietorch import SE3
      2 ex = SE3.Identity(1)
----> 3 ex.quaternion()

~/path/env/lib/python3.8/site-packages/lietorch/groups.py in quaternion(self)
    137     def quaternion(self):
    138         """ extract quaternion """
--> 139         return self.apply_op(Quat, self.data)
    140 
    141     def log(self):

NameError: name 'Quat' is not defined

A quick check reveals that the class Quat, referenced in groups.py:139 is indeed not defined. A convenient method to extract the rotation/quaternion/SO3 object from an SE3 object would be helpful.

zachteed commented 2 years ago

Thanks for bringing this up, I have a more general .vec() implemented but need to add the special case where only the quaternion is extracted.

Right now, you can extract the quaternion using the .vec() function which is implemented for each group. The vec() function is differentiable and extracts the vector embedding for each group. In the case of SE3, it will return a 7d vector with the translation and quaternion [t, q]

from lietorch import SE3
ex = SE3.Identity(1)
t, q = ex.vec().split([3,4], -1)

You can also initialize any group from a vector using the FromVec https://github.com/princeton-vl/lietorch#converting-between-groups-elements-and-euclidean-embeddings