Ros2 node demo3

agimus_controller/agimus_controller/factory/robot_model.py

@@ -11,14 +11,11 @@
 
 @dataclass
 class RobotModelParameters:
-    full_q0: npt.NDArray[np.float64] = field(
-        default_factory=lambda: np.array([], dtype=np.float64)
-    )  # Initial full configuration of the robot
     q0: npt.NDArray[np.float64] = field(
         default_factory=lambda: np.array([], dtype=np.float64)
-    )  # Initial reduced configuration of the robot
+    )  # Initial full configuration of the robot
     free_flyer: bool = False  # True if the robot has a free flyer
-    locked_joint_names: list[str] = field(default_factory=list)
+    moving_joint_names: list[str] = field(default_factory=list)
     urdf: Union[
         Path, str
     ] = ""  # Path to the URDF file or string containing URDF as an XML
@@ -39,26 +36,17 @@ class RobotModelParameters:
     )  # Red color for the collision model
 
     def __post_init__(self):
-        # Check full_q0 is not empty
-        if len(self.full_q0) == 0:
-            raise ValueError("full_q0 cannot be empty.")
-
-        # Check q0 is not empty or if there is no reduced model, q0 is the full configuration
-        if len(self.q0) == 0 and not len(self.locked_joint_names) == 0:
-            raise ValueError("q0 cannot be empty while reducing the model.")
-        elif len(self.q0) == 0:
-            self.q0 = self.full_q0
         # 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)
+            # Use a default armature filled with 0s, based on the size of moving_joint_names
+            self.armature = np.zeros(len(self.moving_joint_names), dtype=np.float64)
 
-        # Ensure armature has the same shape as q0
+        # Ensure armature has the same shape as moving_joint_names
         if (
-            self.armature.shape != self.q0.shape and not self.free_flyer
+            len(self.armature) != len(self.moving_joint_names) and not self.free_flyer
         ):  #! TODO: Do the same for free flyer
             raise ValueError(
-                f"Armature must have the same shape as q0. Got {self.armature.shape} and {self.q0.shape}."
+                f"Armature must have the same shape as moving_joint_names. Got {self.armature.shape} and {len(self.moving_joint_names)}."
             )
 
         # Ensure URDF and SRDF are valid
@@ -96,7 +84,6 @@ def __init__(self, param: RobotModelParameters):
         self._collision_model = None
         self._visual_model = None
         self._q0 = deepcopy(self._params.q0)
-        self._full_q0 = deepcopy(self._params.full_q0)
         self.load_models()  # Populate models
 
     @property
@@ -127,20 +114,10 @@ def collision_model(self) -> pin.GeometryModel:
             raise AttributeError("Colision model has not been computed yet.")
         return self._collision_model
 
-    @property
-    def q0(self) -> np.array:
-        """Initial configuration of the robot."""
-        return self._q0
-
-    @property
-    def full_q0(self) -> np.array:
-        """Initial full configuration of the robot."""
-        return self._full_q0
-
     def load_models(self) -> None:
         """Load and prepare robot models based on parameters."""
         self._load_full_pinocchio_models()
-        self._apply_locked_joints()
+        self._lock_joints()
         if self._params.collision_as_capsule:
             self._update_collision_model_to_capsules()
         if self._params.self_collision:
@@ -187,23 +164,24 @@ def _load_full_pinocchio_models(self) -> None:
                 f"Failed to load URDF models from {self._params.urdf}: {e}"
             )
 
-    def _apply_locked_joints(self) -> None:
+    def _lock_joints(self) -> None:
         """Apply locked joints."""
         joints_to_lock = []
-        for jn in self._params.locked_joint_names:
-            if self._full_robot_model.existJointName(jn):
+        for jn in self._full_robot_model.names:
+            if jn == "universe":
+                continue
+            if not jn in self._params.moving_joint_names:
                 joints_to_lock.append(self._full_robot_model.getJointId(jn))
-            else:
-                raise ValueError(f"Joint {jn} not found in the robot model.")
-
+        if len(self._q0) == 0:
+            self._q0 = np.zeros(self._full_robot_model.nq)
         self._robot_model, [
             self._collision_model,
             self._visual_model,
         ] = pin.buildReducedModel(
             self._full_robot_model,
             [self._collision_model, self._visual_model],
             joints_to_lock,
-            self._full_q0,
+            self._q0,
         )
 
     def _update_collision_model_to_capsules(self) -> None:

agimus_controller/agimus_controller/mpc.py

@@ -26,25 +26,29 @@ def setup(
         self._ocp = ocp
         self._warm_start = warm_start
         self._buffer = buffer
+        self._mpc_debug_data = MPCDebugData(ocp=self._ocp.debug_data)
 
     def run(self, initial_state: TrajectoryPoint, current_time_ns: int) -> OCPResults:
         assert self._ocp is not None
         assert self._warm_start is not None
         timer1 = time.perf_counter_ns()
         self._buffer.clear_past(current_time_ns)
+        if len(self._buffer) < self._ocp.horizon_size:
+            return None
         reference_trajectory = self._extract_horizon_from_buffer()
-        self._ocp.set_reference_trajectory(reference_trajectory)
+        self._ocp.set_reference_weighted_trajectory(reference_trajectory)
         timer2 = time.perf_counter_ns()
+        reference_trajectory_points = [el.point for el in reference_trajectory]
         x0, x_init, u_init = self._warm_start.generate(
-            initial_state, reference_trajectory, self._ocp.debug_data.result
+            initial_state, reference_trajectory_points
         )
         timer3 = time.perf_counter_ns()
         self._ocp.solve(x0, x_init, u_init)
-        self._warm_start.update_previous_solution(self._ocp.debug_data.result)
+        self._warm_start.update_previous_solution(self._ocp._ocp_results.states)
         timer4 = time.perf_counter_ns()
 
         # Extract the solution.
-        self._mpc_debug_data = self._ocp.debug_data
+        self._mpc_debug_data.ocp = self._ocp.debug_data
         self._mpc_debug_data.duration_iteration_ns = timer4 - timer1
         self._mpc_debug_data.duration_horizon_update_ns = timer2 - timer1
         self._mpc_debug_data.duration_generate_warm_start_ns = timer3 - timer2

agimus_controller/agimus_controller/ocp/ocp_croco_goal_reaching.py

@@ -142,8 +142,8 @@ def set_reference_weighted_trajectory(
             # Modify running cost weight
             state_reg.cost.activation.weights = np.concatenate(
                 (
-                    reference_weighted_trajectory[i].weight.w_robot_configuration,
-                    reference_weighted_trajectory[i].weight.w_robot_velocity,
+                    reference_weighted_trajectory[i].weights.w_robot_configuration,
+                    reference_weighted_trajectory[i].weights.w_robot_velocity,
                 )
             )
             # Modify control regularization cost
@@ -155,13 +155,13 @@ def set_reference_weighted_trajectory(
             # Modify running cost weight
             ctrl_reg.cost.activation.weights = reference_weighted_trajectory[
                 i
-            ].weight.w_robot_effort
+            ].weights.w_robot_effort
             # Modify end effector frame cost
 
             # setting running model goal tracking reference, weight and frame id
             # assuming exactly one end-effector tracking reference was passed to the trajectory
             ee_names = list(
-                iter(reference_weighted_trajectory[i].weight.w_end_effector_poses)
+                iter(reference_weighted_trajectory[i].weights.w_end_effector_poses)
             )
             if len(ee_names) > 1:
                 raise ValueError("Only one end-effector tracking reference is allowed.")
@@ -173,7 +173,7 @@ def set_reference_weighted_trajectory(
             ee_cost.cost.residual.id = ee_id
             ee_cost.cost.activation.weights = reference_weighted_trajectory[
                 i
-            ].weight.w_end_effector_poses[ee_name]
+            ].weights.w_end_effector_poses[ee_name]
             ee_cost.cost.residual.reference = reference_weighted_trajectory[
                 i
             ].point.end_effector_poses[ee_name]
@@ -191,14 +191,14 @@ def set_reference_weighted_trajectory(
 
         state_reg.cost.activation.weights = np.concatenate(
             (
-                reference_weighted_trajectory[i].weight.w_robot_configuration,
-                reference_weighted_trajectory[i].weight.w_robot_velocity,
+                reference_weighted_trajectory[-1].weights.w_robot_configuration,
+                reference_weighted_trajectory[-1].weights.w_robot_velocity,
             )
         )
         # Modify end effector frame cost
 
         ee_names = list(
-            iter(reference_weighted_trajectory[i].weight.w_end_effector_poses)
+            iter(reference_weighted_trajectory[-1].weights.w_end_effector_poses)
         )
         ee_name = ee_names[0]
         ee_cost = self._solver.problem.runningModels[-1].differential.costs.costs[
@@ -207,7 +207,7 @@ def set_reference_weighted_trajectory(
         ee_cost.cost.residual.id = ee_id
         ee_cost.cost.activation.weights = reference_weighted_trajectory[
             -1
-        ].weight.w_end_effector_poses[ee_name]
+        ].weights.w_end_effector_poses[ee_name]
         ee_cost.cost.residual.reference = reference_weighted_trajectory[
             -1
         ].point.end_effector_poses[ee_name]

agimus_controller/agimus_controller/ocp_base_croco.py

@@ -103,7 +103,7 @@ def solve(
         self._problem.x0 = x0
         # Solve the OCP
         self._solver.solve(
-            [x0] + x_warmstart, u_warmstart, self._ocp_params.solver_iters
+            list([x0] + x_warmstart), u_warmstart, self._ocp_params.solver_iters
         )
 
         # Store the results

agimus_controller/agimus_controller/trajectory.py

@@ -35,16 +35,17 @@ class WeightedTrajectoryPoint:
     """Trajectory point and it's corresponding weights."""
 
     point: TrajectoryPoint
-    weight: TrajectoryPointWeights
+    weights: TrajectoryPointWeights
 
 
 class TrajectoryBuffer(deque):
     """List of variable size in which the HPP trajectory nodes will be."""
 
-    def clear_past(self, current_time_ns):
-        while self and self[0].point.time_ns < current_time_ns:
-            self.popleft()
+    def clear_past(self, current_time_ns, delta_ns=1e9):
+        self.popleft()
+        # while self and self[0].point.time_ns < current_time_ns - delta_ns:
+        #     self.popleft()
 
-    def horizon(self, horizon_size, dt_ocp):
+    def horizon(self, horizon_size, dt_ocp=None):
         # TBD improve this implementation in case the dt_mpc != dt_ocp
-        return self._buffer[: self._ocp.horizon_size]
+        return list(self)[:horizon_size]

agimus_controller/agimus_controller/warm_start_reference.py

@@ -72,20 +72,22 @@ def generate(
             f"Expected x_init shape {(len(reference_trajectory), self._nx)}, "
             f"from provided reference got {x_init.shape}"
         )
-
+        # print(np.array(reference_trajectory[0].robot_configuration).shape)
+        # print(np.array(reference_trajectory[0].robot_velocity).shape)
+        # print(np.array(reference_trajectory[0].robot_acceleration).shape)
         u_init = [
             pin.rnea(
                 self._rmodel,
                 self._rdata,
-                point.robot_configuration,
-                point.robot_velocity,
-                point.robot_acceleration,
+                np.array(point.robot_configuration),
+                np.array(point.robot_velocity),
+                np.array(point.robot_acceleration),
             )
             for point in reference_trajectory
         ]
         assert np.array(u_init).shape == (len(reference_trajectory), self._rmodel.nv), (
             f"Expected u_init shape {(len(reference_trajectory), self._rmodel.nv)}, "
             f"from provided reference got {u_init.shape}"
         )
-
-        return x0, x_init, u_init
+        # reduce the size of control ref by one to fit Croco way of doing things
+        return x0, x_init, u_init[:-1]

agimus_controller/package.xml

@@ -13,15 +13,7 @@
 
   <!-- Run-time dependencies -->
   <exec_depend>python3-numpy</exec_depend>
-  <exec_depend>colmpc</exec_depend>
-  <exec_depend>crocoddyl</exec_depend>
-  <exec_depend>coal</exec_depend>
-  <exec_depend>example-robot-data</exec_depend>
-  <exec_depend>mim_solvers</exec_depend>
-  <exec_depend>numpy</exec_depend>
-  <exec_depend>pinocchio</exec_depend>
   <exec_depend>xacro</exec_depend>
-  <exec_depend>rospkg</exec_depend>
 
   <!-- Test dependencies -->
   <test_depend>ament_copyright</test_depend>

agimus_controller/tests/test_ocp_croco_base.py

@@ -108,26 +108,28 @@ def set_reference_weighted_trajectory(self, reference_weighted_trajectory):
     def setUpClass(self):
         # Mock the RobotModelFactory and OCPParamsCrocoBase
         robot = robex.load("panda")
-        urdf_path = Path(robot.urdf)
-        srdf_path = Path(robot.urdf.replace("urdf", "srdf"))
-        urdf_meshes_dir = urdf_path.parent.parent.parent.parent.parent
+        urdf = Path(robot.urdf)
+        srdf = Path(robot.urdf.replace("urdf", "srdf"))
+        urdf_meshes_dir = urdf.parent.parent.parent.parent.parent
         free_flyer = False
         locked_joint_names = ["panda_finger_joint1", "panda_finger_joint2"]
         reduced_nq = robot.model.nq - len(locked_joint_names)
-        full_q0 = np.zeros(robot.model.nq)
-        q0 = np.zeros(reduced_nq)
+        moving_joint_names = set(robot.model.names) - set(
+            locked_joint_names + ["universe"]
+        )
+        q0 = np.zeros(robot.model.nq)
         armature = np.full(reduced_nq, 0.1)
+
         # Store shared initial parameters
         self.params = RobotModelParameters(
             q0=q0,
-            full_q0=full_q0,
             free_flyer=free_flyer,
-            locked_joint_names=locked_joint_names,
-            urdf=urdf_path,
-            srdf=srdf_path,
+            moving_joint_names=moving_joint_names,
+            urdf=urdf,
+            srdf=srdf,
             urdf_meshes_dir=urdf_meshes_dir,
             collision_as_capsule=True,
-            self_collision=True,
+            self_collision=False,
             armature=armature,
         )
 

agimus_controller/tests/test_ocp_croco_goal_reaching.py

@@ -24,16 +24,17 @@ def setUp(self):
         free_flyer = False
         locked_joint_names = ["panda_finger_joint1", "panda_finger_joint2"]
         reduced_nq = robot.model.nq - len(locked_joint_names)
-        full_q0 = pin.randomConfiguration(robot.model)
-        q0 = full_q0[:reduced_nq]
+        moving_joint_names = set(robot.model.names) - set(
+            locked_joint_names + ["universe"]
+        )
+        q0 = np.zeros(robot.model.nq)
         armature = np.full(reduced_nq, 0.1)
 
         # Store shared initial parameters
         self.params = RobotModelParameters(
             q0=q0,
-            full_q0=full_q0,
             free_flyer=free_flyer,
-            locked_joint_names=locked_joint_names,
+            moving_joint_names=moving_joint_names,
             urdf=urdf,
             srdf=srdf,
             urdf_meshes_dir=urdf_meshes_dir,