diff --git a/theseus/embodied/collision/collision.py b/theseus/embodied/collision/collision.py index 2e966c22c..d1093199b 100644 --- a/theseus/embodied/collision/collision.py +++ b/theseus/embodied/collision/collision.py @@ -9,7 +9,7 @@ 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 @@ -17,7 +17,7 @@ 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], @@ -25,8 +25,8 @@ def __init__( 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) diff --git a/theseus/embodied/collision/eff_obj_contact.py b/theseus/embodied/collision/eff_obj_contact.py index c164832b9..8ae484313 100644 --- a/theseus/embodied/collision/eff_obj_contact.py +++ b/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]) diff --git a/theseus/embodied/collision/tests/test_collision_factor.py b/theseus/embodied/collision/tests/test_collision_factor.py index 49e21708a..bcc573d80 100644 --- a/theseus/embodied/collision/tests/test_collision_factor.py +++ b/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() diff --git a/theseus/embodied/motionmodel/quasi_static_pushing_planar.py b/theseus/embodied/motionmodel/quasi_static_pushing_planar.py index 3b8cdb877..09abd67bb 100644 --- a/theseus/embodied/motionmodel/quasi_static_pushing_planar.py +++ b/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: diff --git a/theseus/geometry/se2.py b/theseus/geometry/se2.py index 49fda5d61..bb7b63b7e 100644 --- a/theseus/geometry/se2.py +++ b/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:]) diff --git a/theseus/geometry/tests/test_se2.py b/theseus/geometry/tests/test_se2.py index 98a5a0462..e1610498f 100644 --- a/theseus/geometry/tests/test_se2.py +++ b/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]) diff --git a/theseus/utils/examples/tactile_pose_estimation/models.py b/theseus/utils/examples/tactile_pose_estimation/models.py index 47e347e5e..5804ca8bc 100644 --- a/theseus/utils/examples/tactile_pose_estimation/models.py +++ b/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(), ),