From ca35a3cc1baca0b495c73452678ab10c0f7d2013 Mon Sep 17 00:00:00 2001
From: m-guberina <gubi.guberina@gmail.com>
Date: Tue, 4 Mar 2025 11:54:01 +0100
Subject: [PATCH] mobilebaseinterface and mir by itself now make sense and are
 usable for experimentation with navigation. the model property needs to be
 overriden to produce truncated models depending on the control mode to make
 use of base-only ocp for whole body roobots

---
 examples/navigation/mobile_base_navigation.py |  2 -
 python/smc/control/control_loop_manager.py    | 21 ++++----
 .../optimal_control/abstract_croco_ocp.py     |  4 +-
 python/smc/robots/implementations/__init__.py |  2 +-
 python/smc/robots/implementations/heron.py    |  4 +-
 python/smc/robots/implementations/mir.py      | 48 +++++++++++++------
 .../smc/robots/implementations/mobile_yumi.py |  4 ++
 python/smc/robots/implementations/ur5e.py     |  2 +-
 .../interfaces/mobile_base_interface.py       | 28 +++++------
 .../robots/interfaces/whole_body_interface.py | 12 +++++
 python/smc/robots/utils.py                    |  9 +++-
 11 files changed, 89 insertions(+), 47 deletions(-)

diff --git a/examples/navigation/mobile_base_navigation.py b/examples/navigation/mobile_base_navigation.py
index 1dd1127..1020e48 100644
--- a/examples/navigation/mobile_base_navigation.py
+++ b/examples/navigation/mobile_base_navigation.py
@@ -1,5 +1,4 @@
 from smc import getRobotFromArgs
-from smc.robots.interfaces.whole_body_interface import SingleArmWholeBodyInterface
 from smc import getMinimalArgParser
 from smc.path_generation.maps.premade_maps import createSampleStaticMap
 from smc.path_generation.path_math.path2d_to_6d import path2D_to_SE3
@@ -33,7 +32,6 @@ if __name__ == "__main__":
     args = get_args()
     robot = getRobotFromArgs(args)
     # TODO: for ocp you want to pass only the mobile base model
-    # robot._mode = SingleArmWholeBodyInterface.control_mode.base_only
     robot._step()
     robot._q[0] = 9.0
     robot._q[1] = 4.0
diff --git a/python/smc/control/control_loop_manager.py b/python/smc/control/control_loop_manager.py
index 5756795..0e3175d 100644
--- a/python/smc/control/control_loop_manager.py
+++ b/python/smc/control/control_loop_manager.py
@@ -1,5 +1,7 @@
 from argparse import Namespace
 from smc.robots.abstract_robotmanager import AbstractRobotManager
+from smc.robots.interfaces.single_arm_interface import SingleArmInterface
+from smc.robots.interfaces.mobile_base_interface import MobileBaseInterface
 from smc.multiprocessing.process_manager import ProcessManager
 from smc.visualization.plotters import realTimePlotter
 
@@ -149,19 +151,16 @@ class ControlLoopManager:
                         "q": self.robot_manager._q,
                     }
                 )
-                self.robot_manager.visualizer_manager.sendCommand(
-                    {
-                        "T_w_e": self.robot_manager.T_w_e,
-                    }
-                )
-                if self.robot_manager.robot_name == "heron":
-                    T_base = self.robot_manager.data.oMi[1]
+                if issubclass(self.robot_manager.__class__, SingleArmInterface):
+                    self.robot_manager.visualizer_manager.sendCommand(
+                        {
+                            "T_w_e": self.robot_manager.T_w_e,
+                        }
+                    )
+                if issubclass(self.robot_manager.__class__, MobileBaseInterface):
                     self.robot_manager.visualizer_manager.sendCommand(
-                        {"T_base": T_base}
+                        {"T_base": self.robot_manager.T_w_b}
                     )
-                #                if self.robot_manager.robot_name == "yumi":
-                #                    T_base = self.robot_manager.data.oMi[1]
-                #                    self.robot_manager.visualizer_manager.sendCommand({"T_base" : T_base})
 
                 if self.args.visualize_collision_approximation:
                     raise NotImplementedError
diff --git a/python/smc/control/optimal_control/abstract_croco_ocp.py b/python/smc/control/optimal_control/abstract_croco_ocp.py
index 2612724..356a068 100644
--- a/python/smc/control/optimal_control/abstract_croco_ocp.py
+++ b/python/smc/control/optimal_control/abstract_croco_ocp.py
@@ -7,6 +7,8 @@ import importlib.util
 import abc
 from typing import Any
 
+from smc.robots.interfaces.mobile_base_interface import MobileBaseInterface
+
 if importlib.util.find_spec("mim_solvers"):
     try:
         import mim_solvers
@@ -127,7 +129,7 @@ class CrocoOCP(abc.ABC):
 
         # TODO: replace this with subclass or whatever call
         # to check if the robot has mobilebase interface
-        if self.robot.robot_name == "heron" or self.robot.robot_name == "myumi":
+        if issubclass(self.robot.__class__, MobileBaseInterface):
             self.xlb = self.xlb[1:]
             self.xub = self.xub[1:]
 
diff --git a/python/smc/robots/implementations/__init__.py b/python/smc/robots/implementations/__init__.py
index 3871613..1eb77ce 100644
--- a/python/smc/robots/implementations/__init__.py
+++ b/python/smc/robots/implementations/__init__.py
@@ -1,4 +1,4 @@
-# from .mir import AbstractMirRobotManager
+from .mir import SimulatedMirRobotManager
 from .heron import SimulatedHeronRobotManager
 from .mobile_yumi import SimulatedMobileYuMiRobotManager
 from .ur5e import SimulatedUR5eRobotManager
diff --git a/python/smc/robots/implementations/heron.py b/python/smc/robots/implementations/heron.py
index a8b770b..02916b7 100644
--- a/python/smc/robots/implementations/heron.py
+++ b/python/smc/robots/implementations/heron.py
@@ -281,7 +281,9 @@ class RealHeronRobotManager(AbstractHeronRobotManager, AbstractRealRobotManager)
 """
 
 
-def heron_approximation():
+def heron_approximation() -> (
+    tuple[pin.Model, pin.GeometryModel, pin.GeometryModel, pin.Data]
+):
     # arm + gripper
     model_arm, collision_model_arm, visual_model_arm, _ = get_model(
         with_gripper_joints=True
diff --git a/python/smc/robots/implementations/mir.py b/python/smc/robots/implementations/mir.py
index acb22f5..a09e2d9 100644
--- a/python/smc/robots/implementations/mir.py
+++ b/python/smc/robots/implementations/mir.py
@@ -1,11 +1,34 @@
-from smc.robots.abstract_robotmanager import AbstractRobotManager
+from smc.robots.interfaces.mobile_base_interface import MobileBaseInterface
+from smc.robots.abstract_simulated_robotmanager import AbstractSimulatedRobotManager
 
+import pinocchio as pin
+import hppfcl as fcl
+import numpy as np
+from argparse import Namespace
 
-class AbstractMirRobotManager(AbstractRobotManager):
-    pass
+
+class AbstractMirRobotManager(MobileBaseInterface):
+    def __init__(self, args: Namespace):
+        if args.debug_prints:
+            print("AbstractMirRobotManager init")
+        self._model, self._collision_model, self._visual_model, self._data = (
+            mir_approximation()
+        )
+        self._base_frame_id = self.model.getFrameId("mobile_base")
+        self._MAX_ACCELERATION = 1.7  # const
+        self._MAX_QD = 3.14  # const
+        super().__init__(args)
+
+
+class SimulatedMirRobotManager(AbstractSimulatedRobotManager, AbstractMirRobotManager):
+    def __init__(self, args: Namespace):
+        if args.debug_prints:
+            print("SimulatedMirRobotManager init")
+        super().__init__(args)
 
 
 class RealMirRobotManager(AbstractMirRobotManager):
+    # TODO: implement
     def sendVelocityCommandToReal(self, v):
         """
         sendVelocityCommand
@@ -27,17 +50,22 @@ class RealMirRobotManager(AbstractMirRobotManager):
         # self.q = pin.integrate(self.model, self.q, qd * self.dt)
         raise NotImplementedError
 
+    # TODO: implement
     def stopRobot(self):
         raise NotImplementedError
 
+    # TODO: implement
     def setFreedrive(self):
         raise NotImplementedError
 
+    # TODO: implement
     def unSetFreedrive(self):
         raise NotImplementedError
 
 
-def mir_approximation():
+def mir_approximation() -> (
+    tuple[pin.Model, pin.GeometryModel, pin.GeometryModel, pin.Data]
+):
     # mobile base as planar joint (there's probably a better
     # option but whatever right now)
     model_mobile_base = pin.Model()
@@ -47,16 +75,6 @@ def mir_approximation():
     parent_id = 0
     # TEST
     joint_placement = pin.SE3.Identity()
-    # joint_placement.rotation = pin.rpy.rpyToMatrix(0, -np.pi/2, 0)
-    # joint_placement.translation[2] = 0.2
-    # TODO TODO TODO TODO TODO TODO TODO TODO
-    # TODO: heron is actually a differential drive,
-    # meaning that it is not a planar joint.
-    # we could put in a prismatic + revolute joint
-    # as the base (both joints being at same position),
-    # and that should work for our purposes.
-    # this makes sense for initial testing
-    # because mobile yumi's base is a planar joint
     MOBILE_BASE_JOINT_ID = model_mobile_base.addJoint(
         parent_id, pin.JointModelPlanar(), joint_placement.copy(), joint_name
     )
@@ -95,7 +113,7 @@ def mir_approximation():
     # it's tool0 because that's used everywhere
     model_mobile_base.addFrame(
         pin.Frame(
-            "tool0",
+            "mobile_base",
             MOBILE_BASE_JOINT_ID,
             0,
             joint_placement.copy(),
diff --git a/python/smc/robots/implementations/mobile_yumi.py b/python/smc/robots/implementations/mobile_yumi.py
index 52361ea..4aa298e 100644
--- a/python/smc/robots/implementations/mobile_yumi.py
+++ b/python/smc/robots/implementations/mobile_yumi.py
@@ -71,6 +71,10 @@ class RealMobileYumiRobotManager(AbstractMobileYuMiRobotManager):
         self.publisher_joints_cmd.publish(msg)
 
 
+# TODO: define a separate mobile base for YuMi here
+# necessary because the dimensions are not the same
+
+
 def get_mobile_yumi_model() -> (
     tuple[pin.Model, pin.GeometryModel, pin.GeometryModel, pin.Data]
 ):
diff --git a/python/smc/robots/implementations/ur5e.py b/python/smc/robots/implementations/ur5e.py
index af934b8..ef7b7b8 100644
--- a/python/smc/robots/implementations/ur5e.py
+++ b/python/smc/robots/implementations/ur5e.py
@@ -25,7 +25,7 @@ if find_spec("rtde_control"):
 class AbstractUR5eRobotManager(ForceTorqueOnSingleArmWrist):
     def __init__(self, args):
         if args.debug_prints:
-            print("RobotManagerUR5e init")
+            print("AbstractUR5eRobotManager init")
         self._model, self._collision_model, self._visual_model, self._data = get_model()
         self._ee_frame_id = self.model.getFrameId("tool0")
         self._MAX_ACCELERATION = 1.7  # const
diff --git a/python/smc/robots/interfaces/mobile_base_interface.py b/python/smc/robots/interfaces/mobile_base_interface.py
index 56bd0dc..58a05ff 100644
--- a/python/smc/robots/interfaces/mobile_base_interface.py
+++ b/python/smc/robots/interfaces/mobile_base_interface.py
@@ -42,7 +42,7 @@ class MobileBaseInterface(AbstractRobotManager):
     def T_w_b(self):
         return self._T_w_b.copy()
 
-    def computeT_w_b(self, q) -> pin.SE3:
+    def computeT_w_b(self, q: np.ndarray) -> pin.SE3:
         assert type(q) is np.ndarray
         pin.forwardKinematics(
             self.model,
@@ -58,19 +58,19 @@ class MobileBaseInterface(AbstractRobotManager):
     # just driving around with a full-body robot without using hands.
     # not prio right now and I don't want to deal with
     # how this interacts with other interfaces
-    # def forwardKinematics(self):
-    #    pin.forwardKinematics(
-    #        self.model,
-    #        self.data,
-    #        self._q,
-    #    )
-    #    pin.updateFramePlacement(self.model, self.data, self._ee_frame_id)
-    #    self._T_w_e = self.data.oMf[self._ee_frame_id].copy()
-
-    # def _step(self):
-    #    self._updateQ()
-    #    self._updateV()
-    #    self.forwardKinematics()
+    def forwardKinematics(self):
+        pin.forwardKinematics(
+            self.model,
+            self.data,
+            self._q,
+        )
+        pin.updateFramePlacement(self.model, self.data, self._base_frame_id)
+        self._T_w_b = self.data.oMf[self._base_frame_id].copy()
+
+    def _step(self):
+        self._updateQ()
+        self._updateV()
+        self.forwardKinematics()
 
 
 def get_mobile_base_model(underactuated: bool) -> tuple[pin.Model, pin.GeometryModel]:
diff --git a/python/smc/robots/interfaces/whole_body_interface.py b/python/smc/robots/interfaces/whole_body_interface.py
index 4ee9811..d96826b 100644
--- a/python/smc/robots/interfaces/whole_body_interface.py
+++ b/python/smc/robots/interfaces/whole_body_interface.py
@@ -31,6 +31,12 @@ class SingleArmWholeBodyInterface(SingleArmInterface, MobileBaseInterface):
         self._mode: SingleArmWholeBodyInterface.control_mode
         super().__init__(args)
 
+    # TODO: override model property to produce the reduced version
+    # depending on the control mode.
+    # you might want to instantiate all of them in advance for easy switching later
+    # NOTE: that this is currently only important for ocp construction,
+    # even though it's obviously the correct move either way
+
     @property
     def mode(self) -> control_mode:
         return self._mode
@@ -149,6 +155,12 @@ class DualArmWholeBodyInterface(DualArmInterface, MobileBaseInterface):
         self._mode: DualArmWholeBodyInterface.control_mode
         super().__init__(args)
 
+    # TODO: override model property to produce the reduced version
+    # depending on the control mode.
+    # you might want to instantiate all of them in advance for easy switching later
+    # NOTE: that this is currently only important for ocp construction,
+    # even though it's obviously the correct move either way
+
     # TODO: put back in when python3.12 will be safe to use
     #    @override
     @property
diff --git a/python/smc/robots/utils.py b/python/smc/robots/utils.py
index e3875df..c53dda8 100644
--- a/python/smc/robots/utils.py
+++ b/python/smc/robots/utils.py
@@ -28,7 +28,7 @@ def getMinimalArgParser():
         type=str,
         help="which robot you're running or simulating",
         default="ur5e",
-        choices=["ur5e", "heron", "yumi", "myumi"],
+        choices=["ur5e", "heron", "yumi", "myumi", "mir"],
     )
     parser.add_argument(
         "--real",
@@ -275,6 +275,13 @@ def getRobotFromArgs(args: argparse.Namespace) -> AbstractRobotManager:
             # return RealMobileYuMiRobotManager(args)
         else:
             return SimulatedMobileYuMiRobotManager(args)
+    if args.robot == "mir":
+        if args.real:
+            pass
+            # TODO: finish it
+            # return RealMirRobotManager(args)
+        else:
+            return SimulatedMirRobotManager(args)
     raise NotImplementedError(
         f"robot {args.robot} is not supported! run the script you ran with --help to see what's available"
     )
-- 
GitLab