From 87451893fdbc4791246ec7675a8cab129e4fb9da Mon Sep 17 00:00:00 2001
From: m-guberina <gubi.guberina@gmail.com>
Date: Wed, 5 Mar 2025 13:50:32 +0100
Subject: [PATCH] unified control modes, they are also selectable via argument.
 will need to write tests to make sure every case works as expected

---
 development_plan.toml                         | 13 ++++++
 examples/cart_pulling/todos.toml              |  7 +---
 python/smc/robots/abstract_robotmanager.py    | 34 +++++++++++++++
 python/smc/robots/implementations/heron.py    |  5 +--
 .../robots/interfaces/dual_arm_interface.py   | 20 +++------
 .../robots/interfaces/single_arm_interface.py |  3 ++
 .../robots/interfaces/whole_body_interface.py | 41 ++++++++-----------
 python/smc/robots/utils.py                    | 15 +++++++
 8 files changed, 93 insertions(+), 45 deletions(-)

diff --git a/development_plan.toml b/development_plan.toml
index f5abe0a..54761b2 100644
--- a/development_plan.toml
+++ b/development_plan.toml
@@ -222,6 +222,19 @@ purpose = """reducing the amount of code, thereby improving extensibility and ma
 hours_required = 5
 is_done = true
 
+[[project_plan.action_items.action_items]]
+name = "fix (add) --start-from-current-pose feature"
+description = """
+crucial feature to connect simulation and running on the real robot.
+a bit iffy now due to the separation of abstract, real and simulated robot manager. a hack function on the side can get the job done.
+"""
+priority = 1
+deadline = 2025-01-31
+dependencies = []
+purpose = """actually pre-verifying experiments with simulation"""
+hours_required = 2
+is_done = false
+
 
 [[project_plan.action_items]]
 name = "improving extensibility"
diff --git a/examples/cart_pulling/todos.toml b/examples/cart_pulling/todos.toml
index e07e154..3d59d00 100644
--- a/examples/cart_pulling/todos.toml
+++ b/examples/cart_pulling/todos.toml
@@ -181,11 +181,8 @@ we get a bitmap of the map from some ros2 topic. this needs to be starified for
 
 the todos go as follows:
 1) figure out how to get a matrix occupancy grid or whatever format from ros2. 
-2) know which format the path planner wants. it's something in shapely, which is then mapped into star shapes.
-   the outcome here is knowing which shapes are already supported and running with that. 
-   just rectangles and ellipses should be good enough
-3) create a fixed transformation for the static map of the lab into star shapes. ideally
-   save that and keep it cached.
+2) know which format the path planner wants. it's something in shapely, which is then mapped into star shapes. the outcome here is knowing which shapes are already supported and running with that. just rectangles and ellipses should be good enough
+3) create a fixed transformation for the static map of the lab into star shapes. ideally save that and keep it cached.
 4) (optional, but highly preferable) take what isn't static in the occupancy map, convex hull it, hit a rectangle/ellipse on that
 """
 priority = 1
diff --git a/python/smc/robots/abstract_robotmanager.py b/python/smc/robots/abstract_robotmanager.py
index 5a29833..a8da284 100644
--- a/python/smc/robots/abstract_robotmanager.py
+++ b/python/smc/robots/abstract_robotmanager.py
@@ -6,6 +6,7 @@ from argparse import Namespace
 import pinocchio as pin
 import numpy as np
 from functools import partial
+from enum import Enum
 
 
 class AbstractRobotManager(abc.ABC):
@@ -20,11 +21,31 @@ class AbstractRobotManager(abc.ABC):
         - provide functions every robot has to have like getQ
     """
 
+    # NOTE: these are all the possible modes for all possible robots.
+    # all are declared here to enforce uniformity in child classes.
+    # each robot has to declare which of these it supports.
+    control_mode = Enum(
+        "mode",
+        [
+            ("whole_body", 1),
+            ("base_only", 2),
+            ("arm_only", 3),
+            ("both_arms", 3),
+            ("both_arms_only", 3),
+            ("left_arm_only", 4),
+            ("right_arm_only", 5),
+        ],
+    )
+
     def __init__(self, args: Namespace):
         self._model: pin.Model
         self._data: pin.Data
         self._visual_model: pin.pinocchio_pywrap_default.GeometryModel
         self._collision_model: pin.pinocchio_pywrap_default.GeometryModel
+        self._mode: AbstractRobotManager.control_mode = (
+            AbstractRobotManager.control_mode.whole_body
+        )
+        self._available_modes: list[AbstractRobotManager.control_mode]
         self.args: Namespace = args
         if args.debug_prints:
             print("AbstractRobotManager init")
@@ -69,6 +90,10 @@ class AbstractRobotManager(abc.ABC):
         self.gripper = None
         # initialize things that depend on q here
         # self._step()
+        # check for where args.mode is legal is handled by argparse
+        for mode in self.control_mode:
+            if mode.name == args.mode:
+                self.mode = mode
 
         ####################################################################
         #                    processes and utilities robotmanager owns     #
@@ -101,6 +126,15 @@ class AbstractRobotManager(abc.ABC):
                     "sorry, almost done - look at utils to get 90% of the solution!"
                 )
 
+    @property
+    def mode(self) -> control_mode:
+        return self._mode
+
+    @mode.setter
+    def mode(self, mode: control_mode) -> None:
+        assert mode in self._available_modes
+        self._mode = mode
+
     @property
     def model(self) -> pin.Model:
         return self._model
diff --git a/python/smc/robots/implementations/heron.py b/python/smc/robots/implementations/heron.py
index 02916b7..6087d3f 100644
--- a/python/smc/robots/implementations/heron.py
+++ b/python/smc/robots/implementations/heron.py
@@ -29,11 +29,10 @@ class AbstractHeronRobotManager(
 ):
     def __init__(self, args):
         if args.debug_prints:
-            print("RobotManagerHeron init")
+            print("AbstractHeronRobotManager init")
         self._model, self._collision_model, self._visual_model, self._data = (
             heron_approximation()
         )
-        self._mode = SingleArmWholeBodyInterface.control_mode.whole_body
         self._ee_frame_id = self.model.getFrameId("tool0")
         self._base_frame_id = self.model.getFrameId("mobile_base")
         # TODO: CHANGE THIS TO REAL VALUES
@@ -43,7 +42,7 @@ class AbstractHeronRobotManager(
 
 
 class SimulatedHeronRobotManager(
-    AbstractSimulatedRobotManager, AbstractHeronRobotManager
+    AbstractHeronRobotManager, AbstractSimulatedRobotManager
 ):
     def __init__(self, args: Namespace):
         if args.debug_prints:
diff --git a/python/smc/robots/interfaces/dual_arm_interface.py b/python/smc/robots/interfaces/dual_arm_interface.py
index 5bd709a..b04f82a 100644
--- a/python/smc/robots/interfaces/dual_arm_interface.py
+++ b/python/smc/robots/interfaces/dual_arm_interface.py
@@ -1,9 +1,9 @@
 from smc.robots.interfaces.single_arm_interface import SingleArmInterface
+from smc.robots.abstract_robotmanager import AbstractRobotManager
 
 import numpy as np
 import pinocchio as pin
 from argparse import Namespace
-from enum import Enum
 
 
 class DualArmInterface(SingleArmInterface):
@@ -15,10 +15,6 @@ class DualArmInterface(SingleArmInterface):
     - r stands for right
     """
 
-    control_mode = Enum(
-        "mode", [("both_arms", 3), ("left_arm_only", 4), ("right_arm_only", 5)]
-    )
-
     # NOTE: you need to fill in the specific names from the urdf here for your robot
     # (better than putting in a magic number)
     # self._l_ee_frame_id = self.model.getFrameId("left_ee_name")
@@ -34,15 +30,11 @@ class DualArmInterface(SingleArmInterface):
         if args.debug_prints:
             print("DualArmInterface init")
         super().__init__(args)
-
-    @property
-    def mode(self) -> control_mode:
-        return self._mode
-
-    @mode.setter
-    def setMode(self, mode: control_mode) -> None:
-        assert type(mode) is self.control_mode
-        self._mode = mode
+        self._available_modes: list[AbstractRobotManager.control_mode] = [
+            AbstractRobotManager.control_mode.both_arms,
+            AbstractRobotManager.control_mode.left_arm_only,
+            AbstractRobotManager.control_mode.right_arm_only,
+        ]
 
     @property
     def q(self) -> np.ndarray:
diff --git a/python/smc/robots/interfaces/single_arm_interface.py b/python/smc/robots/interfaces/single_arm_interface.py
index 4f03771..f82335b 100644
--- a/python/smc/robots/interfaces/single_arm_interface.py
+++ b/python/smc/robots/interfaces/single_arm_interface.py
@@ -10,6 +10,9 @@ class SingleArmInterface(AbstractRobotManager):
         self._ee_frame_id: int
         self._T_w_e: pin.SE3
         super().__init__(args)
+        self._available_modes: list[AbstractRobotManager.control_mode] = [
+            AbstractRobotManager.control_mode.whole_body,
+        ]
 
     @property
     def ee_frame_id(self):
diff --git a/python/smc/robots/interfaces/whole_body_interface.py b/python/smc/robots/interfaces/whole_body_interface.py
index d96826b..1cc3922 100644
--- a/python/smc/robots/interfaces/whole_body_interface.py
+++ b/python/smc/robots/interfaces/whole_body_interface.py
@@ -1,3 +1,4 @@
+from smc.robots.abstract_robotmanager import AbstractRobotManager
 from smc.robots.interfaces.mobile_base_interface import MobileBaseInterface
 from smc.robots.interfaces.single_arm_interface import SingleArmInterface
 from smc.robots.interfaces.dual_arm_interface import DualArmInterface
@@ -23,12 +24,15 @@ class SingleArmWholeBodyInterface(SingleArmInterface, MobileBaseInterface):
     what you get depends on the mode you set - they're enumerate as above
     """
 
-    control_mode = Enum("mode", [("whole_body", 1), ("base_only", 2), ("arm_only", 3)])
-
     def __init__(self, args: Namespace):
         if args.debug_prints:
             print("SingleArmWholeBodyInterface init")
-        self._mode: SingleArmWholeBodyInterface.control_mode
+        self._mode: AbstractRobotManager.control_mode
+        self._available_modes: list[AbstractRobotManager.control_mode] = [
+            AbstractRobotManager.control_mode.whole_body,
+            AbstractRobotManager.control_mode.base_only,
+            AbstractRobotManager.control_mode.arm_only,
+        ]
         super().__init__(args)
 
     # TODO: override model property to produce the reduced version
@@ -37,14 +41,10 @@ class SingleArmWholeBodyInterface(SingleArmInterface, MobileBaseInterface):
     # 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
-
-    @mode.setter
-    def setMode(self, mode: control_mode) -> None:
-        assert type(mode) is self.control_mode
-        self._mode = mode
+    #    @AbstractRobotManager.mode.setter
+    #    def mode(self, mode: AbstractRobotManager.control_mode) -> None:
+    #        assert type(mode) in self._available_modes
+    #        self._mode = mode
 
     # TODO: put back in when python3.12 will be safe to use
     #    @override
@@ -138,22 +138,17 @@ class SingleArmWholeBodyInterface(SingleArmInterface, MobileBaseInterface):
 
 
 class DualArmWholeBodyInterface(DualArmInterface, MobileBaseInterface):
-    control_mode = Enum(
-        "mode",
-        [
-            ("whole_body", 1),
-            ("base_only", 2),
-            ("both_arms_only", 3),
-            ("left_arm_only", 4),
-            ("right_arm_only", 5),
-        ],
-    )
-
     def __init__(self, args: Namespace):
         if args.debug_prints:
             print("DualArmWholeBodyInterface init")
-        self._mode: DualArmWholeBodyInterface.control_mode
         super().__init__(args)
+        self._available_modes: list[AbstractRobotManager.control_mode] = [
+            AbstractRobotManager.control_mode.whole_body,
+            AbstractRobotManager.control_mode.base_only,
+            AbstractRobotManager.control_mode.both_arms_only,
+            AbstractRobotManager.control_mode.left_arm_only,
+            AbstractRobotManager.control_mode.right_arm_only,
+        ]
 
     # TODO: override model property to produce the reduced version
     # depending on the control mode.
diff --git a/python/smc/robots/utils.py b/python/smc/robots/utils.py
index faf7c5c..95710f6 100644
--- a/python/smc/robots/utils.py
+++ b/python/smc/robots/utils.py
@@ -30,6 +30,21 @@ def getMinimalArgParser():
         default="ur5e",
         choices=["ur5e", "heron", "yumi", "myumi", "mir"],
     )
+    parser.add_argument(
+        "--mode",
+        type=str,
+        help="control mode the robot should run in, i.e. select if you're controlling only a subset of the available joints. most of these only make sense for mobile robots with arms of course (you get an error upon an erroneous selection)",
+        default="whole_body",
+        choices=[
+            "whole_body",
+            "base_only",
+            "arm_only",
+            "both_arms",
+            "both_arms_only",
+            "left_arm_only",
+            "right_arm_only",
+        ],
+    )
     parser.add_argument(
         "--real",
         action=argparse.BooleanOptionalAction,
-- 
GitLab