Add support for SE2 poses in Collision2D (facebookresearch#278)

* Changed SE2.xy() to return a Point2 instead of a tensor.

* Added support for SE2 in Collision2D.

* Minor cleanup in collision jacobians test.

* SE2.translation is now an alias for SE2.xy().

theseus/embodied/collision/collision.py

@@ -9,24 +9,24 @@
 
 from theseus.core import CostFunction, CostWeight, Variable
 from theseus.embodied.kinematics import IdentityModel, KinematicsModel
-from theseus.geometry import Point2
+from theseus.geometry import Point2, SE2
 
 from .signed_distance_field import SignedDistanceField2D
 
 
 class Collision2D(CostFunction):
     def __init__(
         self,
-        pose: Point2,
+        pose: Union[Point2, SE2],
         sdf_origin: Union[Point2, torch.Tensor],
         sdf_data: Union[torch.Tensor, Variable],
         sdf_cell_size: Union[float, torch.Tensor, Variable],
         cost_eps: Union[float, Variable, torch.Tensor],
         cost_weight: CostWeight,
         name: Optional[str] = None,
     ):
-        if not isinstance(pose, Point2):
-            raise ValueError("Collision2D only accepts Point2 poses.")
+        if not isinstance(pose, Point2) and not isinstance(pose, SE2):
+            raise ValueError("Collision2D only accepts Point2 or SE2 poses.")
         super().__init__(cost_weight, name=name)
         self.pose = pose
         self.sdf_origin = SignedDistanceField2D.convert_origin(sdf_origin)
@@ -49,8 +49,19 @@ def __init__(
     def _compute_distances_and_jacobians(
         self,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        robot_state = cast(Point2, self.robot.forward_kinematics(self.pose)["state"])
-        return self.sdf.signed_distance(robot_state.tensor.view(-1, 2, 1))
+        jac_xy: Optional[torch.Tensor] = None
+        if isinstance(self.pose, SE2):
+            aux: List[torch.Tensor] = []
+            xy_pos = self.pose.xy(jacobians=aux)
+            jac_xy = aux[0]
+        else:
+            xy_pos = cast(Point2, self.pose)
+
+        robot_state = self.robot.forward_kinematics(xy_pos)["state"]
+        dist, jac = self.sdf.signed_distance(robot_state.tensor.view(-1, 2, 1))
+        if jac_xy is not None:
+            jac = jac.matmul(jac_xy)
+        return dist, jac
 
     def _error_from_distances(self, distances: torch.Tensor):
         return (self.cost_eps.tensor - distances).clamp(min=0)

theseus/embodied/collision/eff_obj_contact.py

@@ -63,7 +63,7 @@ def _compute_distances_and_jacobians(
         J_transf: List[torch.Tensor] = []
         J_xy: List[torch.Tensor] = []
         eff__obj = self.obj.transform_to(
-            self.eff.xy(jacobians=J_xy), jacobians=J_transf
+            self.eff.xy(jacobians=J_xy).tensor, jacobians=J_transf
         )
         J_transf_obj = J_transf[0]
         J_transf_eff = J_transf[1].matmul(J_xy[0])

theseus/embodied/collision/tests/test_collision_factor.py

@@ -76,13 +76,14 @@ def test_collision2d_copy():
     assert cost_function2.name == "new_name"
 
 
-def test_collision2d_jacobians():
+@pytest.mark.parametrize("pose_cls", [th.Point2, th.SE2])
+def test_collision2d_jacobians(pose_cls):
     rng = torch.Generator()
     rng.manual_seed(0)
     for _ in range(10):
-        for batch_size in [1, 10, 100, 1000]:
+        for batch_size in [1, 10, 100]:
             cost_weight = th.ScaleCostWeight(torch.ones(1).squeeze().double())
-            pose = th.Point2(tensor=torch.randn(batch_size, 2, generator=rng).double())
+            pose = pose_cls.randn(batch_size, generator=rng, dtype=torch.float64)
             origin = th.Point2(torch.ones(batch_size, 2).double())
             sdf_data = th.Variable(
                 torch.randn(batch_size, 10, 10, generator=rng).double()

theseus/embodied/motionmodel/quasi_static_pushing_planar.py

@@ -8,7 +8,7 @@
 import torch
 
 from theseus.core import CostFunction, CostWeight, Variable
-from theseus.geometry import SE2, OptionalJacobians, Point2
+from theseus.geometry import SE2, OptionalJacobians
 from theseus.geometry.so2 import SO2
 
 
@@ -115,7 +115,9 @@ def _compute_V(
         vel_xy_obj__world = self.obj2.xy(jacobians=J_o2xy_o2) - self.obj1.xy(
             jacobians=J_o1xy_o1
         )
-        vel_xy_obj__obj = obj2_angle.unrotate(vel_xy_obj__world, jacobians=J_vxyoo)
+        vel_xy_obj__obj = obj2_angle.unrotate(
+            vel_xy_obj__world.tensor, jacobians=J_vxyoo
+        )
 
         # Putting V together
         obj_diff__world = cast(SE2, self.obj1.between(self.obj2, jacobians=J_odw))
@@ -173,8 +175,8 @@ def _compute_Vp(
         J_vxyco: OptionalJacobians = [] if get_jacobians else None
 
         # Compute contact point velocities using consecutive end effector poses
-        contact_point_1 = self.eff1.xy(jacobians=J_cp1_e1)
-        vel_xy_contact__world = cast(Point2, contact_point_2 - contact_point_1)
+        contact_point_1 = self.eff1.xy(jacobians=J_cp1_e1).tensor
+        vel_xy_contact__world = contact_point_2 - contact_point_1
 
         # Transform velocity to object's 2 frame
         vel_xy_contact__obj = obj2_angle.unrotate(
@@ -220,7 +222,7 @@ def _error_and_jacobians_impl(
         # These quantities are needed by two or more of D, V, Vp, so we compute once.
         obj2_angle = self.obj2.rotation
         self.obj2.theta(jacobians=J_o2angle_o2)
-        contact_point_2 = self.eff2.xy(jacobians=J_cp2_e2)
+        contact_point_2 = self.eff2.xy(jacobians=J_cp2_e2).tensor
 
         # D * V = Vp, See Zhou et al.
         # A Fast Stochastic Contact Model for Planar Pushing and Grasping:

theseus/geometry/se2.py

@@ -134,9 +134,9 @@ def theta(self, jacobians: Optional[List[torch.Tensor]] = None) -> torch.Tensor:
 
     @property
     def translation(self) -> Point2:
-        return Point2(tensor=self[:, :2])
+        return self.xy()
 
-    def xy(self, jacobians: Optional[List[torch.Tensor]] = None) -> torch.Tensor:
+    def xy(self, jacobians: Optional[List[torch.Tensor]] = None) -> Point2:
         if jacobians is not None:
             self._check_jacobians_list(jacobians)
             rotation = self.rotation
@@ -145,7 +145,7 @@ def xy(self, jacobians: Optional[List[torch.Tensor]] = None) -> torch.Tensor:
             )
             J_out[:, :2, :2] = rotation.to_matrix()
             jacobians.append(J_out)
-        return self.translation.tensor
+        return Point2(tensor=self[:, :2])
 
     def update_from_x_y_theta(self, x_y_theta: torch.Tensor):
         rotation = SO2(theta=x_y_theta[:, 2:])

theseus/geometry/tests/test_se2.py

@@ -170,7 +170,7 @@ def test_xy_jacobian(dtype):
         jacobian = []
         se2.xy(jacobians=jacobian)
         expected_jac = numeric_jacobian(
-            lambda groups: th.Point2(groups[0].xy()), [se2], function_dim=2
+            lambda groups: groups[0].xy(), [se2], function_dim=2
         )
         torch.allclose(jacobian[0], expected_jac[0])
 

theseus/utils/examples/tactile_pose_estimation/models.py

@@ -171,7 +171,7 @@ def get_tactile_nn_measurements_inputs(
                 model_measurements.append(
                     torch.cat(
                         (
-                            meas_pose_rel.xy(),
+                            meas_pose_rel.xy().tensor,
                             meas_pose_rel.theta().cos(),
                             meas_pose_rel.theta().sin(),
                         ),