Refactor Robot Model Factory

agimus_controller/agimus_controller/factory/robot_model.py

@@ -1,96 +1,219 @@
 from copy import deepcopy
-from dataclasses import dataclass
-import numpy as np
+from dataclasses import dataclass, field
 from pathlib import Path
+
+import coal
+import numpy as np
+import numpy.typing as npt
 import pinocchio as pin
-from typing import Union
 
 
 @dataclass
 class RobotModelParameters:
-    q0_name = str()
-    free_flyer = False
-    locked_joint_names = []
-    urdf = Path() | str
-    srdf = Path() | str
-    collision_as_capsule = False
-    self_collision = False
+    q0: npt.NDArray[np.float64]  # Initial configuration of the robot
+    free_flyer = False  # True if the robot has a free flyer
+    locked_joint_names: list[str] = field(default_factory=list)
+    urdf_path = Path()  # Path to the URDF file
+    srdf_path = Path()  # Path to the SRDF file
+    urdf_meshes_dir = (
+        Path()
+    )  # Path to the directory containing the meshes and the URDF file.
+    collision_as_capsule = (
+        False  # True if the collision model should be reduced to capsules.
+    )
+    # By default, the collision model when convexified is a sum of spheres and cylinders, often representing capsules. Here, all the couples sphere cylinder sphere are replaced by hppfcl capsules.
+    self_collision = False  # If True, the collision model takes into account collisions pairs written in the srdf file.
+    armature: npt.NDArray[np.float64] = field(
+        default_factory=lambda: np.array([])
+    )  # Default empty NumPy array
+
+    def __post_init__(self):
+        # Check q0 is not empty
+        if len(self.q0) == 0:
+            raise ValueError("q0 cannot be empty.")
+
+        # Handle armature:
+        if self.armature.size == 0:
+            # Use a default armature filled with 0s, based on the size of q0
+            self.armature = np.zeros(len(self.q0), dtype=np.float64)
+
+        # Ensure armature has the same shape as q0
+        if self.armature.shape != self.q0.shape:
+            raise ValueError(
+                f"Armature must have the same shape as q0. Got {self.armature.shape} and {self.q0.shape}."
+            )
+
+        # Ensure paths are valid strings
+        if not isinstance(self.urdf_path, str) or not self.urdf_path:
+            raise ValueError("urdf_path must be a non-empty string.")
+
+        if not isinstance(self.srdf_path, str):
+            raise ValueError("srdf_path must be a string.")
 
 
 class RobotModelFactory:
     """Parse the robot model, reduce it and filter the collision model."""
 
-    """Complete model of the robot."""
-    _complete_model = pin.Model()
-    """ Complete model of the robot with collision meshes. """
-    _complete_collision_model = pin.GeometryModel()
-    """ Complete model of the robot with visualization meshes. """
-    _complete_visual_model = pin.GeometryModel()
-    """ Reduced model of the robot. """
-    _rmodel = pin.Model()
-    """ Reduced model of the robot with visualization meshes. """
-    _rcmodel = pin.GeometryModel()
-    """ Reduced model of the robot with collision meshes. """
-    _rvmodel = pin.GeometryModel()
-    """ Default configuration q0. """
-    _q0 = np.array([])
-    """ Parameters of the model. """
-    _params = RobotModelParameters()
-    """ Path to the collisions environment. """
-    _env = Path()
-
-    def load_model(self, param: RobotModelParameters, env: Union[Path, None]) -> None:
+    def __init__(self, param: RobotModelParameters):
+        """Parse the robot model, reduce it and filter the collision model.
+
+        Args:
+            param (RobotModelParameters): Parameters to load the robot models.
+        """
         self._params = param
-        self._env = env
-        self._load_pinocchio_models(param.urdf, param.free_flyer)
-        self._load_default_configuration(param.srdf, param.q0_name)
-        self._load_reduced_model(param.locked_joint_names, param.q0_name)
-        self._update_collision_model(
-            env, param.collision_as_capsule, param.self_collision, param.srdf
+        self._full_robot_model = None
+        self._robot_model = None
+        self._collision_model = None
+        self._visual_model = None
+        self._q0 = deepcopy(self._params.q0)
+        self.load_models()  # Populate models
+
+    @property
+    def full_robot_model(self) -> pin.Model:
+        """Full robot model."""
+        if self._full_robot_model is None:
+            raise AttributeError("Robot model has not been initialized yet.")
+        return self._full_robot_model
+
+    @property
+    def robot_model(self) -> pin.Model:
+        """Robot model, reduced if specified in the parameters."""
+        if self._robot_model is None:
+            raise AttributeError("Robot model has not been computed yet.")
+        return self._robot_model
+
+    @property
+    def visual_model(self) -> pin.GeometryModel:
+        """Visual model of the robot."""
+        if self._visual_model is None:
+            raise AttributeError("Visual model has not been computed yet.")
+        return self._visual_model
+
+    @property
+    def collision_model(self) -> pin.GeometryModel:
+        """Collision model of the robot."""
+        if self._collision_model is None:
+            raise AttributeError("Visual model has not been computed yet.")
+        return self._collision_model
+
+    @property
+    def q0(self) -> np.array:
+        """Initial configuration of the robot."""
+        return self._q0
+
+    def load_models(self) -> None:
+        """Load and prepare robot models based on parameters."""
+        self._load_full_pinocchio_models()
+        self._apply_locked_joints()
+        self._apply_collision_model_updates()
+
+    def _apply_locked_joints(self) -> None:
+        """Apply locked joints if specified."""
+        if self._params.locked_joint_names:
+            self._load_reduced_model()
+
+    def _apply_collision_model_updates(self) -> None:
+        """Update collision model with specified parameters."""
+        if self._params.collision_as_capsule:
+            self._update_collision_model_to_capsules()
+        if self._params.self_collision:
+            self._update_collision_model_to_self_collision()
+
+    def _get_joints_to_lock(self) -> list[int]:
+        """Get the joints ID to lock.
+
+        Raises:
+            ValueError: Joint name not found in the robot model.
+
+        Returns:
+            list[int]: List of joint IDs to lock.
+        """
+        joints_to_lock = []
+        for jn in self._params.locked_joint_names:
+            if self._full_robot_model.existJointName(jn):
+                joints_to_lock.append(self._full_robot_model.getJointId(jn))
+            else:
+                raise ValueError(f"Joint {jn} not found in the robot model.")
+        return joints_to_lock
+
+    def _load_full_pinocchio_models(self) -> None:
+        """Load the full robot model, the visual model and the collision model."""
+        try:
+            (
+                self._full_robot_model,
+                self._collision_model,
+                self._visual_model,
+            ) = pin.buildModelsFromUrdf(
+                self._params.urdf_path,
+                self._params.urdf_meshes_dir,
+                self._params.free_flyer,
+            )
+        except Exception as e:
+            raise ValueError(
+                f"Failed to load URDF models from {self._params.urdf_path}: {e}"
+            )
+
+    def _load_reduced_model(self) -> None:
+        """Load the reduced robot model."""
+        joints_to_lock = self._get_joints_to_lock()
+        self._robot_model = pin.buildReducedModel(
+            self._full_robot_model, joints_to_lock, self._q0
         )
 
-    def _load_pinocchio_models(self, urdf: Path, free_flyer: bool) -> None:
-        pass
-
-    def _load_default_configuration(self, srdf_path: Path, q0_name: str) -> None:
-        pass
-
-    def _load_reduced_model(self, locked_joint_names, q0_name) -> None:
-        pass
-
-    def _update_collision_model(
-        self,
-        env: Union[Path, None],
-        collision_as_capsule: bool,
-        self_collision: bool,
-        srdf: Path,
-    ) -> None:
-        pass
-
-    def create_complete_robot_model(self) -> pin.Model:
-        return self._complete_model.copy()
-
-    def create_complete_collision_model(self) -> pin.GeometryModel:
-        return self._complete_collision_model.copy()
-
-    def create_complete_visual_model(self) -> pin.GeometryModel:
-        return self._complete_visual_model.copy()
-
-    def create_reduced_robot_model(self) -> pin.Model:
-        return self._rmodel.copy()
-
-    def create_reduced_collision_model(self) -> pin.GeometryModel:
-        return self._rcmodel.copy()
-
-    def create_reduced_visual_model(self) -> pin.GeometryModel:
-        return self._rvmodel.copy()
-
-    def create_default_configuration(self) -> np.array:
-        return self._q0.copy()
-
-    def create_model_parameters(self) -> RobotModelParameters:
-        return deepcopy(self._params)
+    def _update_collision_model_to_capsules(self) -> None:
+        """Update the collision model to capsules."""
+        cmodel = self._collision_model.copy()
+        list_names_capsules = []
+        # Iterate through geometry objects in the collision model
+        for geom_object in cmodel.geometryObjects:
+            geometry = geom_object.geometry
+            # Remove superfluous suffix from the name
+            base_name = "_".join(geom_object.name.split("_")[:-1])
+            # Convert cylinders to capsules
+            if isinstance(geometry, coal.Cylinder):
+                name = self._generate_capsule_name(base_name, list_names_capsules)
+                list_names_capsules.append(name)
+                capsule = pin.GeometryObject(
+                    name=name,
+                    parent_frame=int(geom_object.parentFrame),
+                    parent_joint=int(geom_object.parentJoint),
+                    collision_geometry=coal.Capsule(
+                        geometry.radius, geometry.halfLength
+                    ),
+                    placement=geom_object.placement,
+                )
+                capsule.meshColor = np.array([249, 136, 126, 125]) / 255  # Red color
+                self._collision_model.addGeometryObject(capsule)
+                self._collision_model.removeGeometryObject(geom_object.name)
+
+            # Remove spheres associated with links
+            elif isinstance(geometry, coal.Sphere) and "link" in geom_object.name:
+                self._collision_model.removeGeometryObject(geom_object.name)
+
+    def _update_collision_model_to_self_collision(self) -> None:
+        """Update the collision model to self collision."""
+        pin.addAllCollisionPairs(self._collision_model)
+        pin.removeCollisionPairs(
+            self._robot_model, self._collision_model, self._params.srdf_path
+        )
 
-    def print_model(self):
-        print("full model =\n", self._complete_model)
-        print("reduced model =\n", self._rmodel)
+    def _generate_capsule_name(self, base_name: str, existing_names: list[str]) -> str:
+        """Generates a unique capsule name for a geometry object.
+        Args:
+            base_name (str): The base name of the geometry object.
+            existing_names (list): List of names already assigned to capsules.
+        Returns:
+            str: Unique capsule name.
+        """
+        i = 0
+        while f"{base_name}_capsule_{i}" in existing_names:
+            i += 1
+        return f"{base_name}_capsule_{i}"
+
+    @property
+    def armature(self) -> npt.NDArray[np.float64]:
+        """Armature of the robot.
+        Returns:
+            npt.NDArray[np.float64]: Armature of the robot.
+        """
+        return self._params.armature