Allow constant inputs to cost functions to be passed as floats (faceb…

…ookresearch#150)

* collision and double_integrator variable change

* unit test

theseus/embodied/collision/collision.py

@@ -3,7 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import List, Optional, Tuple, cast
+from typing import List, Optional, Tuple, Union, cast
 
 import torch
 
@@ -22,7 +22,7 @@ def __init__(
         sdf_origin: Variable,
         sdf_data: Variable,
         sdf_cell_size: Variable,
-        cost_eps: Variable,
+        cost_eps: Union[float, Variable, torch.Tensor],
         name: Optional[str] = None,
     ):
         if not isinstance(pose, Point2):
@@ -32,7 +32,13 @@ def __init__(
         self.sdf_origin = sdf_origin
         self.sdf_data = sdf_data
         self.sdf_cell_size = sdf_cell_size
-        self.cost_eps = cost_eps
+        if not isinstance(cost_eps, Variable):
+            if not isinstance(cost_eps, torch.Tensor):
+                cost_eps = torch.tensor(cost_eps)
+            self.cost_eps = Variable(cost_eps)
+        else:
+            self.cost_eps = cost_eps
+        self.cost_eps.data = self.cost_eps.data.view(-1, 1)
         self.register_optim_vars(["pose"])
         self.register_aux_vars(["sdf_origin", "sdf_data", "sdf_cell_size", "cost_eps"])
         self.robot: KinematicsModel = IdentityModel()

theseus/embodied/collision/eff_obj_contact.py

@@ -3,7 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import List, Optional, Tuple, cast
+from typing import List, Optional, Tuple, Union, cast
 
 import torch
 
@@ -23,7 +23,7 @@ def __init__(
         sdf_origin: Variable,
         sdf_data: Variable,
         sdf_cell_size: Variable,
-        eff_radius: Variable,
+        eff_radius: Union[float, Variable, torch.Tensor],
         name: Optional[str] = None,
         use_huber_loss: bool = False,
     ):
@@ -33,7 +33,15 @@ def __init__(
         self.sdf_origin = sdf_origin
         self.sdf_data = sdf_data
         self.sdf_cell_size = sdf_cell_size
-        self.eff_radius = eff_radius
+        if not isinstance(eff_radius, Variable):
+            if not isinstance(eff_radius, torch.Tensor):
+                eff_radius = torch.tensor(eff_radius)
+            self.eff_radius = Variable(eff_radius)
+        else:
+            self.eff_radius = eff_radius
+        if eff_radius.data.squeeze().ndim > 1:
+            raise ValueError("eff_radius must be a 0-D or 1-D tensor.")
+        self.eff_radius.data = self.eff_radius.data.view(-1, 1)
         self.register_optim_vars(["obj", "eff"])
         self.register_aux_vars(
             ["sdf_origin", "sdf_data", "sdf_cell_size", "eff_radius"]

theseus/embodied/collision/tests/test_eff_obj_contact.py

@@ -149,3 +149,42 @@ def test_eff_obj_interesect_errors():
         actual = cost_fn.error()
         expected = outputs["error"][sdf_idx, :]
         assert torch.allclose(actual, expected)
+
+
+def test_eff_obj_variable_type():
+    rng = torch.Generator()
+    rng.manual_seed(0)
+    for batch_size in [1, 10, 100]:
+        obj = create_random_se2(batch_size, rng)
+        eff = create_random_se2(batch_size, rng)
+        origin = th.Variable(torch.randn(batch_size, 2).double())
+        sdf_data = th.Variable(torch.randn(batch_size, 10, 10).double())
+        cell_size = th.Variable(torch.rand(batch_size, 1).double())
+        eff_radius = th.Variable(torch.rand(batch_size, 1).double())
+        cost_weight = th.ScaleCostWeight(1.0)
+        cost_function = th.eb.EffectorObjectContactPlanar(
+            obj, eff, cost_weight, origin, sdf_data, cell_size, eff_radius
+        )
+
+        assert isinstance(cost_function.eff_radius, th.Variable)
+        assert cost_function.eff_radius is eff_radius
+
+        eff_radius_t = torch.rand(batch_size, 1).double()
+
+        cost_function = th.eb.EffectorObjectContactPlanar(
+            obj, eff, cost_weight, origin, sdf_data, cell_size, eff_radius_t
+        )
+
+        assert isinstance(cost_function.eff_radius, th.Variable)
+        assert np.allclose(cost_function.eff_radius.data, eff_radius_t)
+        assert len(cost_function.eff_radius.shape) == 2
+
+        eff_radius_f = torch.rand(1)
+
+        cost_function = th.eb.EffectorObjectContactPlanar(
+            obj, eff, cost_weight, origin, sdf_data, cell_size, eff_radius_f
+        )
+
+        assert isinstance(cost_function.eff_radius, th.Variable)
+        assert np.allclose(cost_function.eff_radius.data.item(), eff_radius_f)
+        assert len(cost_function.eff_radius.shape) == 2

theseus/embodied/motionmodel/double_integrator.py

@@ -18,7 +18,7 @@ def __init__(
         vel1: Vector,
         pose2: LieGroup,
         vel2: Vector,
-        dt: Variable,
+        dt: Union[float, torch.Tensor, Variable],
         cost_weight: CostWeight,
         name: Optional[str] = None,
     ):
@@ -28,16 +28,22 @@ def __init__(
             raise ValueError(
                 "All variables for a DoubleIntegrator must have the same dimension."
             )
+        if not isinstance(dt, Variable):
+            if not isinstance(dt, torch.Tensor):
+                dt = torch.tensor(dt)
+            self.dt = Variable(dt)
+        else:
+            self.dt = dt
         if dt.data.squeeze().ndim > 1:
             raise ValueError(
                 "dt data must be a 0-D or 1-D tensor with numel in {1, batch_size}."
             )
+        self.dt.data = self.dt.data.view(-1, 1)
         self.pose1 = pose1
         self.vel1 = vel1
         self.pose2 = pose2
         self.vel2 = vel2
         self.register_optim_vars(["pose1", "vel1", "pose2", "vel2"])
-        self.dt = dt
         self.register_aux_vars(["dt"])
         self.weight = cost_weight
 
@@ -96,11 +102,13 @@ class GPCostWeight(CostWeight):
     def __init__(
         self,
         Qc_inv: Union[Variable, torch.Tensor],
-        dt: Union[Variable, torch.Tensor],
+        dt: Union[float, Variable, torch.Tensor],
         name: Optional[str] = None,
     ):
         super().__init__(name=name)
         if not isinstance(dt, Variable):
+            if not isinstance(dt, torch.Tensor):
+                dt = torch.tensor(dt)
             dt = Variable(dt)
         if dt.data.squeeze().ndim > 1:
             raise ValueError("dt must be a 0-D or 1-D tensor.")
@@ -172,7 +180,7 @@ def __init__(
         vel1: Vector,
         pose2: LieGroup,
         vel2: Vector,
-        dt: Variable,
+        dt: Union[float, Variable, torch.Tensor],
         cost_weight: GPCostWeight,
         name: Optional[str] = None,
     ):
@@ -181,6 +189,15 @@ def __init__(
                 "GPMotionModel only accepts cost weights of type GPCostWeight. "
                 "For other weight types, consider using DoubleIntegrator instead."
             )
+        if not isinstance(dt, Variable):
+            if not isinstance(dt, torch.Tensor):
+                dt = torch.tensor(dt)
+            self.dt = Variable(dt)
+        else:
+            self.dt = dt
+        if dt.data.squeeze().ndim > 1:
+            raise ValueError("dt must be a 0-D or 1-D tensor.")
+        self.dt.data = self.dt.data.view(-1, 1)
         super().__init__(pose1, vel1, pose2, vel2, dt, cost_weight, name=name)
 
     def _copy_impl(self, new_name: Optional[str] = None) -> "GPMotionModel":

theseus/embodied/motionmodel/tests/test_double_integrator.py

@@ -5,10 +5,12 @@
 
 import copy
 
+import numpy as np
 import pytest  # noqa: F401
 import torch
 
 import theseus as th
+from theseus.core import Variable
 from theseus.core.tests.common import check_another_theseus_function_is_copy
 from theseus.utils import numeric_jacobian
 
@@ -60,6 +62,34 @@ def test_gp_motion_model_cost_weight_copy():
     )
 
 
+def test_gp_motion_model_variable_type():
+    for dof in range(1, 10):
+        for batch_size in [1, 10, 100]:
+            aux = torch.randn(batch_size, dof, dof).double()
+            q_inv = aux.transpose(-2, -1).bmm(aux)
+            dt = torch.rand(1).double()
+            cost_weight = th.eb.GPCostWeight(q_inv, dt)
+
+            assert isinstance(cost_weight.Qc_inv, Variable)
+            assert isinstance(cost_weight.dt, Variable)
+            assert torch.allclose(cost_weight.Qc_inv.data, q_inv)
+            assert torch.allclose(cost_weight.dt.data, dt)
+
+            q_inv_v = Variable(q_inv)
+            dt_v = Variable(dt)
+            cost_weight = th.eb.GPCostWeight(q_inv_v, dt_v)
+            assert isinstance(cost_weight.dt, Variable)
+            assert cost_weight.Qc_inv is q_inv_v
+            assert cost_weight.dt is dt_v
+
+            q_inv_v = Variable(q_inv)
+            dt_f = torch.rand(1)
+            cost_weight = th.eb.GPCostWeight(q_inv_v, dt_f)
+            assert isinstance(cost_weight.dt, Variable)
+            assert np.allclose(cost_weight.dt.data.item(), dt_f)
+            assert len(cost_weight.dt.shape) == 2
+
+
 def test_gp_motion_model_cost_function_error_vector_vars():
     for batch_size in [1, 10, 100]:
         for dof in range(1, 10):