fix inertia computation (#17)

* expand the support for mjcf schema

* inertia fix

* update to robot class

* add more support

* inertia is computed properly

* remove past

* update checks

* pasta againt

* update geometry

* update definition in part

Makefile

@@ -46,11 +46,10 @@ clean:
 #       Static Checks      #
 # ------------------------ #
 
-py-files := $(shell find . -name '*.py')
 
 format:
-	@black $(py-files)
-	@ruff format $(py-files)
+	@black kol
+	@ruff format kol
 .PHONY: format
 
 format-cpp:
@@ -59,9 +58,9 @@ format-cpp:
 .PHONY: format-cpp
 
 static-checks:
-	@black --diff --check $(py-files)
-	@ruff check $(py-files)
-	@mypy --install-types --non-interactive $(py-files)
+	@black --diff --check kol
+	@ruff check kol
+	@mypy --install-types --non-interactive kol
 .PHONY: lint
 
 mypy-daemon:

kol/geometry.py

@@ -25,7 +25,7 @@ def rotation_matrix_to_euler_angles(rotation_matrix: np.matrix) -> tuple[float,
     return x, y, z
 
 
-def apply_matrix_(mesh: stl.mesh.Mesh, matrix: np.matrix) -> stl.mesh.Mesh:
+def apply_matrix_(mesh: stl.mesh.Mesh, matrix: np.ndarray) -> stl.mesh.Mesh:
     rotation = matrix[0:3, 0:3]
     translation = matrix[0:3, 3:4].T.tolist()
 
@@ -43,6 +43,20 @@ def inv_tf(a_to_b_tf: np.matrix) -> np.matrix:
     return np.matrix(np.linalg.inv(a_to_b_tf))
 
 
+def transform_inertia_tensor(inertia: list[float] | np.matrix, rotation: np.ndarray) -> np.ndarray:
+    """Transforms the inertia tensor to a new frame.
+
+    Args:
+        inertia: The inertia tensor in the original frame.
+        rotation: The rotation matrix from the original frame to the new frame.
+
+    Returns:
+        The inertia tensor in the new frame.
+    """
+    inertia_matrix = np.array(inertia).reshape(3, 3)
+    return rotation.T @ inertia_matrix @ rotation
+
+
 @dataclass
 class Dynamics:
     mass: float

kol/onshape/converter.py

@@ -16,9 +16,10 @@
 import networkx as nx
 import numpy as np
 import stl
+from scipy.spatial.transform import Rotation as R
 
 from kol.formats import mjcf, urdf
-from kol.geometry import apply_matrix_, inv_tf, rotation_matrix_to_euler_angles
+from kol.geometry import apply_matrix_, inv_tf, transform_inertia_tensor
 from kol.mesh import MeshExt, stl_to_fmt
 from kol.onshape.api import OnshapeApi
 from kol.onshape.client import OnshapeClient
@@ -137,7 +138,7 @@ def __init__(
         self.cache_map: dict[str, Any] = {}
 
         # Map containing the transformations from the STL origin to the part frame.
-        self.stl_origin_to_part_tf: dict[Key, np.matrix] = {}
+        self.stl_origin_to_part_tf: dict[Key, np.ndarray] = {}
 
     def get_or_use_cached(
         self,
@@ -587,10 +588,10 @@ def get_urdf_part(self, key: Key, joint: Joint | None = None) -> urdf.Link:
         # If the part is the root part, move the STL to be relative to the
         # center of mass and principle inertia axes, otherwise move it to
         # the origin of the part frame.
-        com_to_part_tf = np.matrix(np.eye(4))
-        com_to_part_tf[:3, 3] = -np.array(part_dynamic.center_of_mass).reshape(3, 1)
+        com_to_part_tf = np.eye(4)
+        com_to_part_tf[:3, 3] = -np.array(part_dynamic.center_of_mass).reshape(3)
         if joint is None:
-            stl_origin_to_part_tf = np.matrix(com_to_part_tf)
+            stl_origin_to_part_tf = com_to_part_tf
         else:
             stl_origin_to_part_tf = inv_tf(joint.child_entity.matedCS.part_to_mate_tf)
         self.stl_origin_to_part_tf[key] = stl_origin_to_part_tf
@@ -600,7 +601,6 @@ def get_urdf_part(self, key: Key, joint: Joint | None = None) -> urdf.Link:
 
         if part_file_path.exists():
             logger.info("Using cached file %s", part_file_path)
-
         else:
             # Downloads the STL file.
             part_file_path_stl = part_file_path.with_suffix(".stl")
@@ -609,9 +609,9 @@ def get_urdf_part(self, key: Key, joint: Joint | None = None) -> urdf.Link:
                 buffer = io.BytesIO()
                 self.api.download_stl(part, buffer)
                 buffer.seek(0)
-                mesh = stl.mesh.Mesh.from_file(None, fh=buffer)
-                mesh = apply_matrix_(mesh, stl_origin_to_part_tf)
-                mesh.save(part_file_path_stl)
+                mesh_obj = stl.mesh.Mesh.from_file(None, fh=buffer)
+                mesh_obj = apply_matrix_(mesh_obj, stl_origin_to_part_tf)
+                mesh_obj.save(part_file_path_stl)
 
             # Converts the mesh to the desired format.
             logger.info("Converting STL file to %s", part_file_path)
@@ -626,13 +626,12 @@ def get_urdf_part(self, key: Key, joint: Joint | None = None) -> urdf.Link:
         # joint frame (since URDF expects this by convention).
         mesh_origin = urdf.Origin.zero_origin()
         center_of_mass = part_dynamic.center_of_mass_in_frame(stl_origin_to_part_tf)
-        # inertia = part_dynamic.inertia_in_frame(stl_origin_to_part_tf)
+
         inertia = part_dynamic.inertia_matrix
+        inertia_transformed = transform_inertia_tensor(inertia, stl_origin_to_part_tf[:3, :3])
 
-        # Aligns the inertia matrix with the principal axes.
         principal_axes = part_dynamic.principal_axes_in_frame(stl_origin_to_part_tf)
-        principal_axes_rpy = rotation_matrix_to_euler_angles(principal_axes)
-        # principal_axes_rpy = (0.0, 0.0, 0.0)
+        principal_axes_rpy = R.from_matrix(principal_axes).as_euler("xyz", degrees=False)
 
         urdf_file_path = f"./meshes/{part_file_name}"
         urdf_link_name = self.key_name(key, "link")
@@ -654,12 +653,12 @@ def get_urdf_part(self, key: Key, joint: Joint | None = None) -> urdf.Link:
                 ),
                 mass=mass,
                 inertia=urdf.Inertia(
-                    ixx=float(inertia[0, 0]),
-                    ixy=float(inertia[0, 1]),
-                    ixz=float(inertia[0, 2]),
-                    iyy=float(inertia[1, 1]),
-                    iyz=float(inertia[1, 2]),
-                    izz=float(inertia[2, 2]),
+                    ixx=float(inertia_transformed[0, 0]),
+                    ixy=float(inertia_transformed[0, 1]),
+                    ixz=float(inertia_transformed[0, 2]),
+                    iyy=float(inertia_transformed[1, 1]),
+                    iyz=float(inertia_transformed[1, 2]),
+                    izz=float(inertia_transformed[2, 2]),
                 ),
             ),
             collision=urdf.CollisionLink(
@@ -745,7 +744,7 @@ def resolve(expression: str | None) -> float | None:
                     parent=parent,
                     child=child,
                     origin=origin,
-                    axis=urdf.Axis((0.0, 0.0, -1.0 if joint.lhs_is_first else 1.0)),
+                    axis=urdf.Axis((0.0, 0.0, -1.0)),
                     limits=urdf.JointLimits(
                         effort=effort,
                         velocity=velocity,
@@ -784,7 +783,7 @@ def resolve(expression: str | None) -> float | None:
                     parent=parent,
                     child=child,
                     origin=origin,
-                    axis=urdf.Axis((0.0, 0.0, -1.0 if joint.lhs_is_first else 1.0)),
+                    axis=urdf.Axis((0.0, 0.0, -1.0)),
                     limits=urdf.JointLimits(
                         effort=effort,
                         velocity=velocity,

kol/onshape/schema/part.py

@@ -58,7 +58,7 @@ def principal_inertia(self) -> np.ndarray:
     def principal_axes(self) -> np.matrix:
         return np.matrix(np.array([axis.array for axis in self.principalAxes]))
 
-    def principal_axes_in_frame(self, frame: np.matrix) -> np.matrix:
+    def principal_axes_in_frame(self, frame: np.ndarray) -> np.ndarray:
         return frame[:3, :3] @ self.principal_axes
 
     @property
@@ -73,7 +73,7 @@ def inertia_in_frame(self, frame: np.matrix) -> np.matrix:
     def center_of_mass(self) -> tuple[float, float, float]:
         return (self.centroid[0], self.centroid[1], self.centroid[2])
 
-    def center_of_mass_in_frame(self, frame: np.matrix) -> tuple[float, float, float]:
+    def center_of_mass_in_frame(self, frame: np.ndarray) -> tuple[float, float, float]:
         com = np.matrix(list(self.center_of_mass) + [1.0])
         com_in_frame = (frame * com.T)[:3]
         return (float(com_in_frame[0, 0]), float(com_in_frame[1, 0]), float(com_in_frame[2, 0]))

kol/requirements.txt

@@ -7,6 +7,7 @@ numpy
 numpy-stl
 pydantic
 mujoco
+scipy
 
 # Types
 types-requests

pyproject.toml

@@ -37,6 +37,7 @@ module = [
     "networkx.*",
     "pybullet.*",
     "stl.*",
+    "scipy.*",
 ]
 
 ignore_missing_imports = true