Fixed some mypy issues. Added drm to dev requirements so unit tests p…

…ass.

examples/homography_estimation.py

@@ -17,7 +17,7 @@
 from torch.utils.data import DataLoader, Dataset
 
 import theseus as th
-from theseus.third_party.easyaug import RandomGeoAug, GeoAugParam, RandomPhotoAug
+from theseus.third_party.easyaug import GeoAugParam, RandomGeoAug, RandomPhotoAug
 from theseus.third_party.utils import grid_sample
 
 FONT = cv2.FONT_HERSHEY_DUPLEX
@@ -36,10 +36,10 @@ def prepare_data():
     dataset_root = os.path.join(os.getcwd(), "data")
     chunks = [
         "revisitop1m.1",
-        #"revisitop1m.2", # Uncomment for more data.
-        #"revisitop1m.3",
-        #"revisitop1m.4",
-        #"revisitop1m.5",
+        # "revisitop1m.2", # Uncomment for more data.
+        # "revisitop1m.3",
+        # "revisitop1m.4",
+        # "revisitop1m.5",
     ]
     dataset_paths = []
     for chunk in chunks:

requirements/dev.txt

@@ -4,6 +4,7 @@ mypy>=0.910
 nox==2020.8.22
 pre-commit>=2.9.2
 isort>=5.6.4
+differentiable-robot-model>=0.2.3
 types-PyYAML==5.4.3
 numdifftools>=0.9.40
 mock>=4.0.3

theseus/embodied/kinematics/kinematics_model.py

@@ -40,8 +40,10 @@ def __init__(self, urdf_path: str):
         try:
             import differentiable_robot_model as drm
         except ModuleNotFoundError as e:
-            print("UrdfRobotModel requires installing differentiable-robot-model.")
-            print("Please run `pip install differentiable-robot-model`.")
+            print(
+                "UrdfRobotModel requires installing differentiable-robot-model. "
+                "Please run `pip install differentiable-robot-model`."
+            )
             raise e
 
         self.drm_model = drm.DifferentiableRobotModel(urdf_path)

theseus/optimizer/linear_system.py

@@ -16,7 +16,7 @@ def __init__(
         row_ptr: np.ndarray,
         num_rows: int,
         num_cols: int,
-        dtype: np.dtype = np.float_,
+        dtype: np.dtype = np.float_,  # type: ignore
     ):
         self.col_ind = col_ind
         self.row_ptr = row_ptr

theseus/theseus_layer.py

@@ -298,7 +298,7 @@ def jacobians(self) -> Tuple[List[torch.Tensor], torch.Tensor]:
         return [self.var.project(euclidean_grad, is_sparse=True)], self.error()
 
     def dim(self) -> int:
-        return np.prod(self.var.tensor.shape[1:])
+        return int(np.prod(self.var.tensor.shape[1:]))
 
     def _copy_impl(self, new_name: Optional[str] = None) -> "CostFunction":
         return _DLMPerturbation(

theseus/utils/examples/tactile_pose_estimation/misc.py

@@ -47,11 +47,11 @@ def __init__(
             stop = max(int(np.ceil(num_episodes * val_ratio)), 2)
             idx = order[:stop] if data_mode == "val" else order[stop:]
 
-        self.img_feats = data["img_feats"][idx]
-        self.eff_poses = data["eff_poses"][idx]
-        self.obj_poses = data["obj_poses"][idx]
-        self.contact_episode = data["contact_episode"][idx]
-        self.contact_flag = data["contact_flag"][idx]
+        self.img_feats = data["img_feats"][idx]  # type: ignore
+        self.eff_poses = data["eff_poses"][idx]  # type: ignore
+        self.obj_poses = data["obj_poses"][idx]  # type: ignore
+        self.contact_episode = data["contact_episode"][idx]  # type: ignore
+        self.contact_flag = data["contact_flag"][idx]  # type: ignore
         # Check sizes of the attributes assigned above
         self.dataset_size: int = -1
         for key in data: