From 10b35b1fb9b5d35f17d7c5ad8c88dc01c4a3be65 Mon Sep 17 00:00:00 2001
From: m-guberina <gubi.guberina@gmail.com>
Date: Wed, 5 Mar 2025 21:10:52 +0100
Subject: [PATCH] mobile navigation with drawing works, now it's super easy to
 test navigation by itself

---
 development_plan.toml                         |  12 ++
 .../mpc_base_clik_single_arm_control_loop.py  |  14 +--
 examples/navigation/mobile_base_navigation.py | 102 ++++++++++++++---
 .../path_following_template.py                |  34 +++---
 .../controller_templates/point_to_point.py    |  51 +++++++++
 .../croco_mpc_path_following.py               |  67 ++++++-----
 .../croco_mpc_point_to_point.py               |  54 +++++++++
 .../point_to_point_croco_ocp.py               | 108 +++++++++---------
 python/smc/robots/abstract_robotmanager.py    |   2 -
 python/smc/robots/implementations/mir.py      |   2 +-
 python/smc/robots/interfaces/__init__.py      |   5 +
 .../robots/interfaces/dual_arm_interface.py   |   1 +
 .../interfaces/mobile_base_interface.py       |   5 +
 .../robots/interfaces/single_arm_interface.py |   1 +
 python/smc/util/draw_path.py                  |   3 +-
 .../meshcat_viewer_wrapper/visualizer.py      |   5 +-
 python/smc/visualization/visualizer.py        |   5 +-
 17 files changed, 332 insertions(+), 139 deletions(-)
 create mode 100644 python/smc/robots/interfaces/__init__.py

diff --git a/development_plan.toml b/development_plan.toml
index 54761b2..d1c9084 100644
--- a/development_plan.toml
+++ b/development_plan.toml
@@ -235,6 +235,18 @@ purpose = """actually pre-verifying experiments with simulation"""
 hours_required = 2
 is_done = false
 
+[[project_plan.action_items.action_items]]
+name = "refactor croco mpc p2p"
+description = """
+you have 4 or 5 literally identical controlloops and it looks (is) stupid. the circle to square is that they need to take in different arguments due to the templates. find how to intelligently package this.
+"""
+priority = 1
+deadline = 2025-01-31
+dependencies = []
+purpose = """actually pre-verifying experiments with simulation"""
+hours_required = 3
+is_done = false
+
 
 [[project_plan.action_items]]
 name = "improving extensibility"
diff --git a/examples/cart_pulling/disjoint_control/mpc_base_clik_single_arm_control_loop.py b/examples/cart_pulling/disjoint_control/mpc_base_clik_single_arm_control_loop.py
index 6f2d7e9..2e5d09d 100644
--- a/examples/cart_pulling/disjoint_control/mpc_base_clik_single_arm_control_loop.py
+++ b/examples/cart_pulling/disjoint_control/mpc_base_clik_single_arm_control_loop.py
@@ -13,7 +13,7 @@ from smc.path_generation.path_math.cart_pulling_path_math import (
 )
 from smc.path_generation.path_math.path_to_trajectory import path2D_to_trajectory2D
 from smc.control.controller_templates.path_following_template import (
-    PathFollowingFromPlannerControlLoop,
+    PathFollowingFromPlannerCtrllLoopTemplate,
 )
 from smc.control.cartesian_space.ik_solvers import dampedPseudoinverse
 from smc.control.optimal_control.croco_mpc_path_following import initializePastData
@@ -48,7 +48,7 @@ def BaseMPCEECLIKPathFollowingFromPlannerMPCControlLoop(
     v_cmd = xs[1, robot.model.nq :]
 
     pathSE3_handlebar = construct_EE_path(args, p, past_data["path2D_untimed"])
-    robot._mode = SingleArmWholeBodyInterface.control_mode.arm_only
+    robot._mode = SingleArmWholeBodyInterface.control_mode.upper_body
 
     T_w_e = robot.T_w_e
     # first check whether we're at the goal
@@ -62,10 +62,10 @@ def BaseMPCEECLIKPathFollowingFromPlannerMPCControlLoop(
 
     v_cmd[3:] = v_arm
 
-    if args.visualizer:
-        if t % args.viz_update_rate == 0:
-            robot.visualizer_manager.sendCommand({"path": path_base})
-            robot.visualizer_manager.sendCommand({"frame_path": pathSE3_handlebar})
+    #    if args.visualizer:
+    #        if t % args.viz_update_rate == 0:
+    #            robot.visualizer_manager.sendCommand({"path": path_base})
+    #            robot.visualizer_manager.sendCommand({"frame_path": pathSE3_handlebar})
 
     err_vector_ee = log6(robot.T_w_e.actInv(pathSE3_handlebar[0]))
     err_vector_base = np.linalg.norm(p - path_base[0][:2])  # z axis is irrelevant
@@ -104,7 +104,7 @@ def BaseMPCANDEECLIKCartPulling(
     else:
         get_position = lambda robot: robot.q[:2]
         controlLoop = partial(
-            PathFollowingFromPlannerControlLoop,
+            PathFollowingFromPlannerCtrllLoopTemplate,
             path_planner,
             get_position,
             ocp,
diff --git a/examples/navigation/mobile_base_navigation.py b/examples/navigation/mobile_base_navigation.py
index 1020e48..d218389 100644
--- a/examples/navigation/mobile_base_navigation.py
+++ b/examples/navigation/mobile_base_navigation.py
@@ -1,59 +1,123 @@
 from smc import getRobotFromArgs
 from smc import getMinimalArgParser
+from smc.robots.interfaces import MobileBaseInterface
 from smc.path_generation.maps.premade_maps import createSampleStaticMap
 from smc.path_generation.path_math.path2d_to_6d import path2D_to_SE3
 from smc.control.optimal_control.util import get_OCP_args
 from smc.control.cartesian_space import getClikArgs
 from smc.path_generation.planner import starPlanner, getPlanningArgs
+from smc.control.optimal_control.croco_mpc_point_to_point import CrocoBaseP2PMPC
 from smc.control.optimal_control.croco_mpc_path_following import (
     CrocoBasePathFollowingMPC,
 )
 from smc.multiprocessing import ProcessManager
+from smc.util.draw_path import drawPath
 
-import time
 import numpy as np
 from functools import partial
 import pinocchio as pin
+import argparse
+import matplotlib
 
+# TODO:
+# try valueing future points on the path more the current ones!!!!
 
-def get_args():
+
+def get_args() -> argparse.Namespace:
     parser = getMinimalArgParser()
     parser = get_OCP_args(parser)
     parser = getClikArgs(parser)  # literally just for goal error
     parser = getPlanningArgs(parser)
+    parser.add_argument(
+        "--planner",
+        action=argparse.BooleanOptionalAction,
+        help="if on, you're in a pre-set map and a planner produce a plan to navigate. if off, you draw the path to be followed",
+        default=True,
+    )
+    parser.add_argument(
+        "--draw-new",
+        action=argparse.BooleanOptionalAction,
+        help="are you drawing a new path or reusing the previous one",
+        default=False,
+    )
+    parser.add_argument(
+        "--map-width",
+        type=float,
+        help="width of the map in meters (x-axis) - only used for drawing of the path",
+        default=3.0,
+    )
+    parser.add_argument(
+        "--map-height",
+        type=float,
+        help="height of the map in meters (y-axis) - only used for drawing of the path",
+        default=3.0,
+    )
     args = parser.parse_args()
     return args
 
 
-# DA BI OVAJ MPC SEX FUNKCIONIRAO, MORAS MU DAT PIN.MODEL SAMO BAZE!!!
-# ili to il ce se vrtit GRO sporo bez razloga
+def fixedPathPlanner(path2D: np.ndarray, p_base: np.ndarray) -> np.ndarray:
+    """
+    fixedPathPlanner
+    ----------------
+    we can assume robot is following this path and that it started the following controller at the first point.
+    to make the controller make sense (without a sense of timing), we need to find the closest point on the path,
+    and delete the already traversed portion of the path.
+    the assumption is that the controller constructs a trajectory from the path, starting with the first point.
+
+    the cleanest way to get this done is associate robot's current position with the path parameter s,
+    but i'll just implement it as quickly as possible, not in the most correct or the most efficient manner.
+
+    NOTE: path is (N,3) in shape, last column is 0, because subotimal code structure
+    """
+
+    p_base_3 = np.array([p_base[0], p_base[1], 0.0])
+    distances = np.linalg.norm(p_base_3 - path2D, axis=1)
+    index = np.argmin(distances)
+    path2D = path2D[index:]
+    return path2D
+
 
 if __name__ == "__main__":
     args = get_args()
     robot = getRobotFromArgs(args)
-    # TODO: for ocp you want to pass only the mobile base model
-    robot._step()
-    robot._q[0] = 9.0
-    robot._q[1] = 4.0
+    assert issubclass(robot.__class__, MobileBaseInterface)
     robot._step()
+    if args.planner:
+        robot._q[0] = 9.0
+        robot._q[1] = 4.0
+        robot._step()
     x0 = np.concatenate([robot.q, robot.v])
     goal = np.array([0.5, 5.5])
     T_w_b = robot.T_w_b
 
-    planning_function = partial(starPlanner, goal)
-    # here we're following T_w_e reference so that's what we send
-    path_planner = ProcessManager(
-        args, planning_function, T_w_b.translation[:2], 3, None
-    )
-    _, map_as_list = createSampleStaticMap()
-    if args.visualizer:
-        robot.sendRectangular2DMapToVisualizer(map_as_list)
-        # time.sleep(5)
+    if args.planner:
+        planning_function = partial(starPlanner, goal)
+        # here we're following T_w_e reference so that's what we send
+        path_planner = ProcessManager(
+            args, planning_function, T_w_b.translation[:2], 3, None
+        )
+        _, map_as_list = createSampleStaticMap()
+        if args.visualizer:
+            robot.sendRectangular2DMapToVisualizer(map_as_list)
+            # time.sleep(5)
+    else:
+        if not args.draw_new:
+            pixel_path_file_path = "./parameters/path_in_pixels.csv"
+            path2D = np.genfromtxt(pixel_path_file_path, delimiter=",")
+        else:
+            matplotlib.use("tkagg")
+            path2D = drawPath(args)
+            matplotlib.use("qtagg")
+        path2D[:, 0] = path2D[:, 0] * args.map_width
+        path2D[:, 1] = path2D[:, 1] * args.map_height
+        path2D = np.hstack((path2D, np.zeros((len(path2D), 1))))
+        robot.updateViz({"fixed_path": path2D})
+        CrocoBaseP2PMPC(args, robot, path2D[0])
+        path_planner = partial(fixedPathPlanner, path2D)
 
     CrocoBasePathFollowingMPC(args, robot, x0, path_planner)
 
-    print("final position:", robot.T_w_b.translation)
-
     if args.real:
         robot.stopRobot()
 
diff --git a/python/smc/control/controller_templates/path_following_template.py b/python/smc/control/controller_templates/path_following_template.py
index b33d4ec..1ff64ef 100644
--- a/python/smc/control/controller_templates/path_following_template.py
+++ b/python/smc/control/controller_templates/path_following_template.py
@@ -1,3 +1,4 @@
+from types import FunctionType
 from smc.robots.abstract_robotmanager import AbstractRobotManager
 from smc.multiprocessing.process_manager import ProcessManager
 
@@ -10,8 +11,8 @@ global control_loop_return
 control_loop_return = tuple[bool, dict[str, np.ndarray], dict[str, np.ndarray]]
 
 
-def PathFollowingFromPlannerControlLoop(
-    path_planner: ProcessManager,
+def PathFollowingFromPlannerCtrllLoopTemplate(
+    path_planner: ProcessManager|FunctionType,
     get_position: Callable[[AbstractRobotManager], np.ndarray],
     SOLVER: Any,
     control_loop: Callable[
@@ -42,23 +43,28 @@ def PathFollowingFromPlannerControlLoop(
 
     p = get_position(robot)
 
-    path_planner.sendCommand(p)
-    data = path_planner.getData()
+    if type(path_planner) == ProcessManager:
+        path_planner.sendCommand(p)
+        data = path_planner.getData()
 
-    if data == "done":
-        breakFlag = True
+        if data == "done":
+            breakFlag = True
 
-    if data == "done" or data is None:
-        robot.sendVelocityCommand(np.zeros(robot.model.nv))
-        log_item["qs"] = robot.q
-        log_item["dqs"] = robot.v
-        return breakFlag, save_past_item, log_item
+        if data == "done" or data is None:
+            robot.sendVelocityCommand(np.zeros(robot.model.nv))
+            log_item["qs"] = robot.q
+            log_item["dqs"] = robot.v
+            return breakFlag, save_past_item, log_item
 
-    _, path2D_untimed = data
-    path2D_untimed = np.array(path2D_untimed).reshape((-1, 2))
+        _, path2D = data
+        path2D = np.array(path2D).reshape((-1, 2))
+    else:
+        path2D = path_planner(p)
+        if len(path2D) < 4:
+            breakFlag = True
 
     v_cmd, past_item_inner, log_item_inner = control_loop(
-        SOLVER, path2D_untimed, args, robot, t, past_data
+        SOLVER, path2D, args, robot, t, past_data
     )
     robot.sendVelocityCommand(v_cmd)
     log_item.update(log_item_inner)
diff --git a/python/smc/control/controller_templates/point_to_point.py b/python/smc/control/controller_templates/point_to_point.py
index 113f5e0..c4937fd 100644
--- a/python/smc/control/controller_templates/point_to_point.py
+++ b/python/smc/control/controller_templates/point_to_point.py
@@ -1,3 +1,4 @@
+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
 
@@ -15,6 +16,56 @@ from smc.robots.interfaces.whole_body_interface import (
 global control_loop_return
 control_loop_return = tuple[bool, dict[str, np.ndarray], dict[str, np.ndarray]]
 
+def BaseP2PCtrlLoopTemplate(
+    SOLVER: Any,
+    p_basegoal: np.ndarray,
+    control_loop: Callable[
+        [
+            Any,
+            np.ndarray,
+            Namespace,
+            MobileBaseInterface,
+            int,
+            dict[str, deque[np.ndarray]],
+        ],
+        tuple[np.ndarray, dict[str, np.ndarray], dict[str, np.ndarray]],
+    ],
+    args: Namespace,
+    robot: MobileBaseInterface,
+    t: int,  # will be float eventually
+    past_data: dict[str, deque[np.ndarray]],
+) -> control_loop_return:
+    """
+    EEAndBaseP2PCtrlLoopTemplate
+    ---------------
+    generic control loop for point to point motion for end-effectors of a dual arm robot
+    (handling error to final point etc).
+    """
+    breakFlag = False
+    log_item = {}
+    save_past_item = {}
+
+    v_cmd, past_item_inner, log_item_inner = control_loop(
+        SOLVER, p_basegoal, args, robot, t, past_data
+    )
+    robot.sendVelocityCommand(v_cmd)
+    log_item.update(log_item_inner)
+    save_past_item.update(past_item_inner)
+
+    p_base = np.array(list(robot.q[:2]) + [0.0])
+    base_err_vector_norm = np.linalg.norm(p_basegoal - p_base)
+
+    if  (
+        base_err_vector_norm < robot.args.goal_error
+    ):
+        breakFlag = True
+
+    log_item["qs"] = robot.q
+    log_item["dqs"] = robot.v
+    log_item["dqs_cmd"] = v_cmd.reshape((robot.model.nv,))
+    log_item["base_err_norm"] = base_err_vector_norm.reshape((1,))
+    return breakFlag, save_past_item, log_item
+
 
 def EEP2PCtrlLoopTemplate(
     SOLVER: Any,
diff --git a/python/smc/control/optimal_control/croco_mpc_path_following.py b/python/smc/control/optimal_control/croco_mpc_path_following.py
index 7efc59a..6a98cdc 100644
--- a/python/smc/control/optimal_control/croco_mpc_path_following.py
+++ b/python/smc/control/optimal_control/croco_mpc_path_following.py
@@ -19,7 +19,7 @@ from smc.path_generation.path_math.cart_pulling_path_math import (
 )
 from smc.path_generation.path_math.path_to_trajectory import path2D_to_trajectory2D
 from smc.control.controller_templates.path_following_template import (
-    PathFollowingFromPlannerControlLoop,
+    PathFollowingFromPlannerCtrllLoopTemplate,
 )
 
 import numpy as np
@@ -32,7 +32,7 @@ from collections import deque
 
 def CrocoBasePathFollowingFromPlannerMPCControlLoop(
     ocp: CrocoOCP,
-    path2D_untimed: np.ndarray,
+    path2D: np.ndarray,
     args: Namespace,
     robot: MobileBaseInterface,
     t: int,
@@ -41,7 +41,7 @@ def CrocoBasePathFollowingFromPlannerMPCControlLoop(
 
     p = robot.T_w_b.translation[:2]
     max_base_v = np.linalg.norm(robot._max_v[:2])
-    path_base = path2D_to_trajectory2D(args, path2D_untimed, max_base_v)
+    path_base = path2D_to_trajectory2D(args, path2D, max_base_v)
     path_base = np.hstack((path_base, np.zeros((len(path_base), 1))))
 
     if args.visualizer:
@@ -74,19 +74,16 @@ def CrocoBasePathFollowingMPC(
     x0 = np.concatenate([robot.q, robot.v])
     ocp.solveInitialOCP(x0)
 
-    if type(path_planner) == types.FunctionType:
-        raise NotImplementedError
-    else:
-        get_position = lambda robot: robot.T_w_b.translation[:2]
-        controlLoop = partial(
-            PathFollowingFromPlannerControlLoop,
-            path_planner,
-            get_position,
-            ocp,
-            CrocoBasePathFollowingFromPlannerMPCControlLoop,
-            args,
-            robot,
-        )
+    get_position = lambda robot: robot.T_w_b.translation[:2]
+    controlLoop = partial(
+        PathFollowingFromPlannerCtrllLoopTemplate,
+        path_planner,
+        get_position,
+        ocp,
+        CrocoBasePathFollowingFromPlannerMPCControlLoop,
+        args,
+        robot,
+    )
     log_item = {
         "qs": np.zeros(robot.nq),
         "dqs": np.zeros(robot.nv),
@@ -147,7 +144,7 @@ def CrocoEEPathFollowingMPCControlLoop(
 
 def CrocoEEPathFollowingFromPlannerMPCControlLoop(
     ocp: CrocoOCP,
-    path2D_untimed: np.ndarray,
+    path2D: np.ndarray,
     args: Namespace,
     robot: SingleArmInterface,
     t: int,
@@ -165,7 +162,7 @@ def CrocoEEPathFollowingFromPlannerMPCControlLoop(
     max_base_v = np.linalg.norm(robot._max_v[:2])
     perc_of_max_v = 0.5
     velocity = perc_of_max_v * max_base_v
-    trajectory2D = path2D_to_trajectory2D(args, path2D_untimed, velocity)
+    trajectory2D = path2D_to_trajectory2D(args, path2D, velocity)
 
     # create a 3D reference out of the path
     path_EE_SE3 = path2D_to_SE3(trajectory2D, args.handlebar_height)
@@ -215,7 +212,7 @@ def CrocoEEPathFollowingMPC(
     else:
         get_position = lambda robot: robot.T_w_e.translation[:2]
         controlLoop = partial(
-            PathFollowingFromPlannerControlLoop,
+            PathFollowingFromPlannerCtrllLoopTemplate,
             path_planner,
             get_position,
             ocp,
@@ -284,7 +281,7 @@ def BaseAndEEPathFollowingMPCControlLoop(
 
 def BaseAndEEPathFollowingFromPlannerMPCControlLoop(
     ocp: CrocoOCP,
-    path2D_untimed_base: np.ndarray,
+    path2D_base: np.ndarray,
     args: Namespace,
     robot: SingleArmWholeBodyInterface,
     t: int,
@@ -294,10 +291,10 @@ def BaseAndEEPathFollowingFromPlannerMPCControlLoop(
 
     # NOTE: this one is obtained as the future path from path planner
     max_base_v = np.linalg.norm(robot._max_v[:2])
-    trajectory_base = path2D_to_trajectory2D(args, path2D_untimed_base, max_base_v)
+    trajectory_base = path2D_to_trajectory2D(args, path2D_base, max_base_v)
     trajectory_base = np.hstack((trajectory_base, np.zeros((len(trajectory_base), 1))))
 
-    trajectorySE3_handlebar = construct_EE_path(args, p, past_data["path2D_untimed"])
+    trajectorySE3_handlebar = construct_EE_path(args, p, past_data["path2D"])
 
     if args.visualizer:
         if t % int(np.ceil(args.ctrl_freq / 25)) == 0:
@@ -315,7 +312,7 @@ def BaseAndEEPathFollowingFromPlannerMPCControlLoop(
     log_item["err_vec_ee"] = err_vector_ee
     log_item["err_norm_ee"] = np.linalg.norm(err_vector_ee).reshape((1,))
     log_item["err_norm_base"] = np.linalg.norm(err_vector_base).reshape((1,))
-    save_past_item = {"path2D_untimed": p}
+    save_past_item = {"path2D": p}
     return v_cmd, save_past_item, log_item
 
 
@@ -362,7 +359,7 @@ def BaseAndEEPathFollowingMPC(
     else:
         get_position = lambda robot: robot.q[:2]
         controlLoop = partial(
-            PathFollowingFromPlannerControlLoop,
+            PathFollowingFromPlannerCtrllLoopTemplate,
             path_planner,
             get_position,
             ocp,
@@ -377,14 +374,14 @@ def BaseAndEEPathFollowingMPC(
         "err_norm_ee": np.zeros((1,)),
         "err_norm_base": np.zeros((1,)),
     }
-    save_past_dict = {"path2D_untimed": T_w_e.translation[:2]}
+    save_past_dict = {"path2D": T_w_e.translation[:2]}
     loop_manager = ControlLoopManager(
         robot, controlLoop, args, save_past_dict, log_item
     )
 
     # actually put past data into the past window
-    loop_manager.past_data["path2D_untimed"].clear()
-    loop_manager.past_data["path2D_untimed"].extend(
+    loop_manager.past_data["path2D"].clear()
+    loop_manager.past_data["path2D"].extend(
         path2D_handlebar[i] for i in range(args.past_window_size)
     )
 
@@ -397,7 +394,7 @@ def BaseAndDualEEPathFollowingMPCControlLoop(
     T_absgoal_l: SE3,
     T_absgoal_r: SE3,
     ocp: CrocoOCP,
-    path2D_untimed_base: np.ndarray,
+    path2D_base: np.ndarray,
     args: Namespace,
     robot: DualArmWholeBodyInterface,
     t: int,
@@ -412,10 +409,10 @@ def BaseAndDualEEPathFollowingMPCControlLoop(
 
     # NOTE: this one is obtained as the future path from path planner
     max_base_v = np.linalg.norm(robot._max_v[:2])
-    trajectory_base = path2D_to_trajectory2D(args, path2D_untimed_base, max_base_v)
+    trajectory_base = path2D_to_trajectory2D(args, path2D_base, max_base_v)
     trajectory_base = np.hstack((trajectory_base, np.zeros((len(trajectory_base), 1))))
 
-    trajectorySE3_handlebar = construct_EE_path(args, p, past_data["path2D_untimed"])
+    trajectorySE3_handlebar = construct_EE_path(args, p, past_data["path2D"])
     trajectorySE3_l = []
     trajectorySE3_r = []
     for traj_pose in trajectorySE3_handlebar:
@@ -443,7 +440,7 @@ def BaseAndDualEEPathFollowingMPCControlLoop(
     log_item["err_vec_ee_r"] = err_vector_ee_r
     log_item["err_norm_ee_r"] = err_norm_ee_r.reshape((1,))
     log_item["err_norm_base"] = np.linalg.norm(err_vector_base).reshape((1,))
-    save_past_item = {"path2D_untimed": p}
+    save_past_item = {"path2D": p}
     return v_cmd, save_past_item, log_item
 
 def BaseAndDualEEPathFollowingMPC(
@@ -480,7 +477,7 @@ def BaseAndDualEEPathFollowingMPC(
         get_position = lambda robot: robot.q[:2]
         BaseAndDualEEPathFollowingMPCControlLoop_with_l_r = partial(BaseAndDualEEPathFollowingMPCControlLoop, T_absgoal_l, T_absgoal_r)
         controlLoop = partial(
-            PathFollowingFromPlannerControlLoop,
+            PathFollowingFromPlannerCtrllLoopTemplate,
             path_planner,
             get_position,
             ocp,
@@ -498,13 +495,13 @@ def BaseAndDualEEPathFollowingMPC(
         "err_norm_ee_r": np.zeros((1,)),
         "err_norm_base": np.zeros((1,)),
     }
-    save_past_dict = {"path2D_untimed": T_w_abs.translation[:2]}
+    save_past_dict = {"path2D": T_w_abs.translation[:2]}
     loop_manager = ControlLoopManager(
         robot, controlLoop, args, save_past_dict, log_item
     )
     # actually put past data into the past window
-    loop_manager.past_data["path2D_untimed"].clear()
-    loop_manager.past_data["path2D_untimed"].extend(
+    loop_manager.past_data["path2D"].clear()
+    loop_manager.past_data["path2D"].extend(
         path2D_handlebar[i] for i in range(args.past_window_size)
     )
 
diff --git a/python/smc/control/optimal_control/croco_mpc_point_to_point.py b/python/smc/control/optimal_control/croco_mpc_point_to_point.py
index a2cdd9f..061ea50 100644
--- a/python/smc/control/optimal_control/croco_mpc_point_to_point.py
+++ b/python/smc/control/optimal_control/croco_mpc_point_to_point.py
@@ -1,9 +1,11 @@
 from smc.control.controller_templates.point_to_point import (
+    BaseP2PCtrlLoopTemplate,
     DualEEAndBaseP2PCtrlLoopTemplate,
     DualEEP2PCtrlLoopTemplate,
     EEAndBaseP2PCtrlLoopTemplate,
     EEP2PCtrlLoopTemplate,
 )
+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
 from smc.robots.interfaces.whole_body_interface import (
@@ -12,6 +14,7 @@ from smc.robots.interfaces.whole_body_interface import (
 )
 from smc.control.optimal_control.abstract_croco_ocp import CrocoOCP
 from smc.control.optimal_control.point_to_point_croco_ocp import (
+    BaseIKOCP,
     SingleArmIKOCP,
     DualArmIKOCP,
     BaseAndSingleArmIKOCP,
@@ -25,6 +28,57 @@ from functools import partial
 from collections import deque
 from argparse import Namespace
 
+def CrocoBaseP2PMPCControlLoop(
+    ocp: CrocoOCP,
+    p_basegoal: np.ndarray,
+    args: Namespace,
+    robot: MobileBaseInterface,
+    _: int,
+    __: dict[str, deque[np.ndarray]],
+) -> tuple[np.ndarray, dict[str, np.ndarray], dict[str, np.ndarray]]:
+    """
+    CrocoIKMPCControlLoop
+    ---------------------
+    """
+    # set initial state from sensor
+    x0 = np.concatenate([robot.q, robot.v])
+    ocp.warmstartAndReSolve(x0)
+    xs = np.array(ocp.solver.xs)
+    # NOTE: for some reason the first value is always some wild bs
+    vel_cmd = xs[1, robot.model.nq :]
+    return vel_cmd, {}, {}
+
+
+def CrocoBaseP2PMPC(
+    args: Namespace, robot: MobileBaseInterface, p_basegoal: np.ndarray, run=True
+):
+    """
+    CrocoBaseP2PMPC
+    -----
+    base point-to-point task
+    """
+    x0 = np.concatenate([robot.q, robot.v])
+    ocp = BaseIKOCP(args, robot, x0, p_basegoal)
+    ocp.solveInitialOCP(x0)
+
+    controlLoop = partial(
+        BaseP2PCtrlLoopTemplate, ocp, p_basegoal, CrocoBaseP2PMPCControlLoop, args, robot
+    )
+    log_item = {
+        "qs": np.zeros(robot.nq),
+        "base_err_norm": np.zeros(1),
+        "dqs": np.zeros(robot.nv),
+        "dqs_cmd": np.zeros(robot.nv),
+    }
+    save_past_dict = {}
+    loop_manager = ControlLoopManager(
+        robot, controlLoop, args, save_past_dict, log_item
+    )
+    if run:
+        loop_manager.run()
+    else:
+        return loop_manager
+
 
 def CrocoEEP2PMPCControlLoop(
     ocp: CrocoOCP,
diff --git a/python/smc/control/optimal_control/point_to_point_croco_ocp.py b/python/smc/control/optimal_control/point_to_point_croco_ocp.py
index 519458f..5fec443 100644
--- a/python/smc/control/optimal_control/point_to_point_croco_ocp.py
+++ b/python/smc/control/optimal_control/point_to_point_croco_ocp.py
@@ -11,6 +11,51 @@ import pinocchio as pin
 import crocoddyl
 from argparse import Namespace
 
+class BaseIKOCP(CrocoOCP):
+    def __init__(
+        self,
+        args: Namespace,
+        robot: SingleArmInterface,
+        x0: np.ndarray,
+        p_basegoal: pin.SE3,
+    ):
+        super().__init__(args, robot, x0, p_basegoal)
+
+    def constructTaskObjectiveFunction(self, goal) -> None:
+        p_basegoal = goal
+        for i in range(self.args.n_knots):
+            baseTranslationResidual = crocoddyl.ResidualModelFrameTranslation(
+                self.state, self.robot.base_frame_id, p_basegoal, self.state.nv
+            )
+            baseTrackingCost = crocoddyl.CostModelResidual(
+                self.state, baseTranslationResidual
+            )
+            self.runningCostModels[i].addCost(
+                "base_translation" + str(i),
+                baseTrackingCost,
+                self.args.base_translation_cost,
+            )
+        self.terminalCostModel.addCost(
+            "base_translation" + str(self.args.n_knots),
+            baseTrackingCost,
+            self.args.base_translation_cost,
+        )
+
+    # there is nothing to update in a point-to-point task
+    def updateCosts(self, data):
+        pass
+
+    def updateGoalInModels(self, goal) -> None:
+        p_basegoal = goal
+        for i, runningModel in enumerate(self.solver.problem.runningModels):
+            runningModel.differential.costs.costs[
+                "base_translation" + str(i)
+            ].cost.residual.reference = p_basegoal
+        self.solver.problem.terminalModel.differential.costs.costs[
+            "base_translation" + str(self.args.n_knots)
+        ].cost.residual.reference = p_basegoal
+
+
 
 class SingleArmIKOCP(CrocoOCP):
     def __init__(
@@ -152,43 +197,23 @@ class DualArmIKOCP(CrocoOCP):
         ].cost.residual.reference = T_w_rgoal
 
 
-class BaseAndSingleArmIKOCP(SingleArmIKOCP):
+class BaseAndSingleArmIKOCP(SingleArmIKOCP, BaseIKOCP):
     def __init__(
         self,
         args: Namespace,
         robot: SingleArmWholeBodyInterface,
         x0: np.ndarray,
         goal,
-        #        T_w_eegoal: pin.SE3,
-        #        p_basegoal: np.ndarray,
     ):
-        # T_w_eegoal, p_basegoal = goal
         super().__init__(args, robot, x0, goal)
 
     def constructTaskObjectiveFunction(
-        # self, T_w_eegoal: pin.SE3, p_basegoal: np.ndarray
         self,
         goal,
     ) -> None:
         T_w_eegoal, p_basegoal = goal
         super().constructTaskObjectiveFunction(T_w_eegoal)
-        for i in range(self.args.n_knots):
-            baseTranslationResidual = crocoddyl.ResidualModelFrameTranslation(
-                self.state, self.robot.base_frame_id, p_basegoal, self.state.nv
-            )
-            baseTrackingCost = crocoddyl.CostModelResidual(
-                self.state, baseTranslationResidual
-            )
-            self.runningCostModels[i].addCost(
-                "base_translation" + str(i),
-                baseTrackingCost,
-                self.args.base_translation_cost,
-            )
-        self.terminalCostModel.addCost(
-            "base_translation" + str(self.args.n_knots),
-            baseTrackingCost,
-            self.args.base_translation_cost,
-        )
+        BaseIKOCP.constructTaskObjectiveFunction(self, p_basegoal)
 
     # there is nothing to update in a point-to-point task
     def updateCosts(self, data):
@@ -198,16 +223,10 @@ class BaseAndSingleArmIKOCP(SingleArmIKOCP):
         # self, T_w_eegoal: pin.SE3, p_basegoal: np.ndarray) -> None:
         T_w_eegoal, p_basegoal = goal
         super().updateGoalInModels(T_w_eegoal)
-        for i, runningModel in enumerate(self.solver.problem.runningModels):
-            runningModel.differential.costs.costs[
-                "base_translation" + str(i)
-            ].cost.residual.reference = p_basegoal
-        self.solver.problem.terminalModel.differential.costs.costs[
-            "base_translation" + str(self.args.n_knots)
-        ].cost.residual.reference = p_basegoal
+        BaseIKOCP.updateGoalInModels(self, p_basegoal)
 
 
-class BaseAndDualArmIKOCP(DualArmIKOCP):
+class BaseAndDualArmIKOCP(DualArmIKOCP, BaseIKOCP):
     def __init__(
         self,
         args: Namespace,
@@ -215,46 +234,21 @@ class BaseAndDualArmIKOCP(DualArmIKOCP):
         x0: np.ndarray,
         goal,
     ):
-        # T_w_eegoal, p_basegoal = goal
         super().__init__(args, robot, x0, goal)
 
     def constructTaskObjectiveFunction(
-        # self, T_w_eegoal: pin.SE3, p_basegoal: np.ndarray
         self,
         goal,
     ) -> None:
         T_w_lgoal, T_w_rgoal, p_basegoal = goal
         super().constructTaskObjectiveFunction((T_w_lgoal, T_w_rgoal))
-        for i in range(self.args.n_knots):
-            baseTranslationResidual = crocoddyl.ResidualModelFrameTranslation(
-                self.state, self.robot.base_frame_id, p_basegoal, self.state.nv
-            )
-            baseTrackingCost = crocoddyl.CostModelResidual(
-                self.state, baseTranslationResidual
-            )
-            self.runningCostModels[i].addCost(
-                "base_translation" + str(i),
-                baseTrackingCost,
-                self.args.base_translation_cost,
-            )
-        self.terminalCostModel.addCost(
-            "base_translation" + str(self.args.n_knots),
-            baseTrackingCost,
-            self.args.base_translation_cost,
-        )
+        BaseIKOCP.constructTaskObjectiveFunction(self, p_basegoal)
 
     # there is nothing to update in a point-to-point task
     def updateCosts(self, data) -> None:
         pass
 
     def updateGoalInModels(self, goal) -> None:
-        # self, T_w_eegoal: pin.SE3, p_basegoal: np.ndarray) -> None:
         T_w_lgoal, T_w_rgoal, p_basegoal = goal
         super().updateGoalInModels(T_w_lgoal, T_w_rgoal)
-        for i, runningModel in enumerate(self.solver.problem.runningModels):
-            runningModel.differential.costs.costs[
-                "base_translation" + str(i)
-            ].cost.residual.reference = p_basegoal
-        self.solver.problem.terminalModel.differential.costs.costs[
-            "base_translation" + str(self.args.n_knots)
-        ].cost.residual.reference = p_basegoal
+        BaseIKOCP.updateGoalInModels(self, p_basegoal)
diff --git a/python/smc/robots/abstract_robotmanager.py b/python/smc/robots/abstract_robotmanager.py
index 10bec21..5be02d2 100644
--- a/python/smc/robots/abstract_robotmanager.py
+++ b/python/smc/robots/abstract_robotmanager.py
@@ -130,8 +130,6 @@ class AbstractRobotManager(abc.ABC):
 
     @mode.setter
     def mode(self, mode: control_mode) -> None:
-        print(mode)
-        print(self._available_modes)
         assert mode in self._available_modes
         self._mode = mode
 
diff --git a/python/smc/robots/implementations/mir.py b/python/smc/robots/implementations/mir.py
index a09e2d9..7bc8742 100644
--- a/python/smc/robots/implementations/mir.py
+++ b/python/smc/robots/implementations/mir.py
@@ -106,7 +106,7 @@ def mir_approximation() -> (
         "box_shape", MOBILE_BASE_JOINT_ID, box_shape, body_placement.copy()
     )
 
-    geometry_mobile_base.meshColor = np.array([1.0, 0.1, 0.1, 1.0])
+    geometry_mobile_base.meshColor = np.array([1.0, 0.1, 0.1, 0.3])
     geom_model_mobile_base.addGeometryObject(geometry_mobile_base)
 
     # have to add the frame manually
diff --git a/python/smc/robots/interfaces/__init__.py b/python/smc/robots/interfaces/__init__.py
new file mode 100644
index 0000000..bc7ff83
--- /dev/null
+++ b/python/smc/robots/interfaces/__init__.py
@@ -0,0 +1,5 @@
+from .dual_arm_interface import *
+from .force_torque_sensor_interface import *
+from .mobile_base_interface import *
+from .single_arm_interface import *
+from .whole_body_interface import *
diff --git a/python/smc/robots/interfaces/dual_arm_interface.py b/python/smc/robots/interfaces/dual_arm_interface.py
index 98045ce..a19fbd0 100644
--- a/python/smc/robots/interfaces/dual_arm_interface.py
+++ b/python/smc/robots/interfaces/dual_arm_interface.py
@@ -34,6 +34,7 @@ class DualArmInterface(SingleArmInterface):
         if not hasattr(self, "_available_modes"):
             self._available_modes: list[AbstractRobotManager.control_mode] = [
                 AbstractRobotManager.control_mode.whole_body,
+                AbstractRobotManager.control_mode.upper_body,  # in this case the same as wholebody
                 AbstractRobotManager.control_mode.left_arm_only,
                 AbstractRobotManager.control_mode.right_arm_only,
             ]
diff --git a/python/smc/robots/interfaces/mobile_base_interface.py b/python/smc/robots/interfaces/mobile_base_interface.py
index 58a05ff..022502b 100644
--- a/python/smc/robots/interfaces/mobile_base_interface.py
+++ b/python/smc/robots/interfaces/mobile_base_interface.py
@@ -25,6 +25,11 @@ class MobileBaseInterface(AbstractRobotManager):
             print("MobileBase init")
         self._base_frame_id: int
         self._T_w_b: pin.SE3
+        if not hasattr(self, "_available_modes"):
+            self._available_modes: list[AbstractRobotManager.control_mode] = [
+                AbstractRobotManager.control_mode.whole_body,
+                AbstractRobotManager.control_mode.base_only,  # in this case the same as wholebody
+            ]
         super().__init__(args)
 
     @property
diff --git a/python/smc/robots/interfaces/single_arm_interface.py b/python/smc/robots/interfaces/single_arm_interface.py
index 667cd3d..31f0045 100644
--- a/python/smc/robots/interfaces/single_arm_interface.py
+++ b/python/smc/robots/interfaces/single_arm_interface.py
@@ -14,6 +14,7 @@ class SingleArmInterface(AbstractRobotManager):
         if not hasattr(self, "_available_modes"):
             self._available_modes: list[AbstractRobotManager.control_mode] = [
                 AbstractRobotManager.control_mode.whole_body,
+                AbstractRobotManager.control_mode.upper_body,  # in this case the same as wholebody
             ]
         super().__init__(args)
 
diff --git a/python/smc/util/draw_path.py b/python/smc/util/draw_path.py
index 9e89708..3b3a7f2 100644
--- a/python/smc/util/draw_path.py
+++ b/python/smc/util/draw_path.py
@@ -10,6 +10,7 @@ possible improvement: make it all bezier curves
       (but if that was a parameter that would be ok i guess)
 """
 
+from argparse import Namespace
 import numpy as np
 import matplotlib.pyplot as plt
 
@@ -56,7 +57,7 @@ class DrawPathManager:
             plt.close()
 
 
-def drawPath(args):
+def drawPath(args:Namespace)->np.ndarray:
     # normalize both x and y to 0-1 range
     # we can multiply however we want later
     # idk about the number of points, but it's large enough to draw
diff --git a/python/smc/visualization/meshcat_viewer_wrapper/visualizer.py b/python/smc/visualization/meshcat_viewer_wrapper/visualizer.py
index ca2e1fe..9a506ea 100644
--- a/python/smc/visualization/meshcat_viewer_wrapper/visualizer.py
+++ b/python/smc/visualization/meshcat_viewer_wrapper/visualizer.py
@@ -107,11 +107,12 @@ class MeshcatVisualizer(PMV):
         self.n_points += 1
 
     def addPath(
-        self, name, path: list[pin.SE3] | np.ndarray, radius=5e-3, color=[1, 0, 0, 0.8]
+        self, name, path: list[pin.SE3] | np.ndarray, radius=5e-3, color=[0, 0, 1, 0.8]
     ):
         # who cares about the name
-        name = "path"
         self.n_path_points = len(path)
+        if name == "fixed_path":
+            color=[1, 0, 0, 0.8]
         if type(path) == np.ndarray:
             # complete the quartenion
             path = np.hstack((path, np.zeros((len(path), 3))))
diff --git a/python/smc/visualization/visualizer.py b/python/smc/visualization/visualizer.py
index 041dcba..3abc3df 100644
--- a/python/smc/visualization/visualizer.py
+++ b/python/smc/visualization/visualizer.py
@@ -71,7 +71,10 @@ def manipulatorVisualizer(
                     viz.addBoxObstacle(cmd["obstacle_box"][0], cmd["obstacle_box"][1])
                 if key == "path":
                     # stupid and evil but there is no time
-                    viz.addPath("", cmd["path"])
+                    viz.addPath("path", cmd["path"])
+                if key == "fixed_path":
+                    # stupid and evil but there is no time
+                    viz.addPath("fixed_path", cmd["fixed_path"])
                 if key == "frame_path":
                     # stupid and evil but there is no time
                     viz.addFramePath("", cmd["frame_path"])
-- 
GitLab