Add warmstart implementation for warmstart from reference points

agimus_controller/agimus_controller/trajectory.py

@@ -1,32 +1,33 @@
 from collections import deque
 from dataclasses import dataclass
 import numpy as np
+import numpy.typing as npt
 from pinocchio import SE3, Force
 
 
 @dataclass
 class TrajectoryPoint:
     """Trajectory point aiming at being a reference for the MPC."""
 
-    time_ns: int
-    robot_configuration: np.ndarray
-    robot_velocity: np.ndarray
-    robot_acceleration: np.ndarray
-    robot_effort: np.ndarray
-    forces: dict[Force]  # Dictionary of pinocchio.Force
-    end_effector_poses: dict[SE3]  # Dictionary of pinocchio.SE3
+    time_ns: int | None = None
+    robot_configuration: npt.NDArray[np.float64] | None = None
+    robot_velocity: npt.NDArray[np.float64] | None = None
+    robot_acceleration: npt.NDArray[np.float64] | None = None
+    robot_effort: npt.NDArray[np.float64] | None = None
+    forces: dict[Force] | None = None  # Dictionary of pinocchio.Force
+    end_effector_poses: dict[SE3] | None = None  # Dictionary of pinocchio.SE3
 
 
 @dataclass
 class TrajectoryPointWeights:
     """Trajectory point weights aiming at being set in the MPC costs."""
 
-    w_robot_configuration: np.ndarray
-    w_robot_velocity: np.ndarray
-    w_robot_acceleration: np.ndarray
-    w_robot_effort: np.ndarray
-    w_forces: dict[np.ndarray]
-    w_end_effector_poses: dict[np.ndarray]
+    w_robot_configuration: npt.NDArray[np.float64] | None = None
+    w_robot_velocity: npt.NDArray[np.float64] | None = None
+    w_robot_acceleration: npt.NDArray[np.float64] | None = None
+    w_robot_effort: npt.NDArray[np.float64] | None = None
+    w_forces: dict[npt.NDArray[np.float64]] | None = None
+    w_end_effector_poses: dict[npt.NDArray[np.float64]] | None = None
 
 
 @dataclass

agimus_controller/agimus_controller/warm_start_base.py

@@ -1,24 +1,69 @@
 from abc import ABC, abstractmethod
+
 import numpy as np
+import numpy.typing as npt
 
 from agimus_controller.trajectory import TrajectoryPoint
 
 
 class WarmStartBase(ABC):
+    """
+    A template class for implementing methods that generate a warmstart for an optimal control problem.
+    The warmstart is expected to generate the initial values for state and control trajectories,
+    based on the initial robot state and a reference trajectory.
+
+    Attributes:
+        _previous_solution (list[TrajectoryPoint]): Stores the previous solution of the optimization problem.
+
+    Methods:
+        generate(initial_state: TrajectoryPoint, reference_trajectory: list[TrajectoryPoint]) -> tuple:
+            Generates warm-start values for the optimization problem. This must be
+            implemented by subclasses.
+
+        update_previous_solution(previous_solution: list[TrajectoryPoint]) -> None:
+            Updates the internally stored previous solution for later use.
+    """
+
     def __init__(self) -> None:
         super().__init__()
         self._previous_solution: list[TrajectoryPoint] = list()
 
     @abstractmethod
     def generate(
         self,
+        initial_state: TrajectoryPoint,
         reference_trajectory: list[TrajectoryPoint],
-    ) -> tuple(np.ndarray, np.ndarray):
-        """Returns x_init, u_init."""
+    ) -> tuple[
+        npt.NDArray[np.float64],
+        list[npt.NDArray[np.float64]],
+        list[npt.NDArray[np.float64]],
+    ]:
+        """
+        Generate initial values for a warm-start of the optimization problem.
+
+        Args:
+            initial_state (TrajectoryPoint): The initial state of the robot,
+                containing `robot_configuration` and `robot_velocity`.
+            reference_trajectory (list[TrajectoryPoint]): A list of `TrajectoryPoint` objects
+                representing the reference trajectory.
+
+        Returns:
+            tuple:
+                - x0 (npt.NDArray[np.float64]): The initial state vector.
+                - init_xs (list[npt.NDArray[np.float64]]): List of state vectors
+                for each point in the reference trajectory.
+                - init_us (list[npt.NDArray[np.float64]]): List of control inputs.
+        """
         ...
 
     def update_previous_solution(
         self, previous_solution: list[TrajectoryPoint]
     ) -> None:
-        """Stores internally the previous solution of the OCP"""
+        """
+        Store the solution from a previous optimization cycle.
+        It can be used as a reference or initialization of warmstart.
+
+        Args:
+            previous_solution (list[TrajectoryPoint]): The solution of the optimization problem.
+        """
         self._previous_solution = previous_solution

agimus_controller/agimus_controller/warm_start_reference.py

@@ -0,0 +1,122 @@
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import override
+
+import pinocchio as pin
+
+from agimus_controller.warm_start_base import WarmStartBase
+from agimus_controller.trajectory import TrajectoryPoint
+
+
+class WarmStartReference(WarmStartBase):
+    """
+    A class for generating warmstart values for trajectory optimization problem.
+
+    This class uses a reference trajectory and the robot model to compute the initial state,
+    state vectors, and control inputs.
+
+    Attributes:
+    _rmodel (pin.Model | None): The robot's Pinocchio model, used for forward dynamics computations.
+    _rdata (pin.Data | None): Data structure associated with the Pinocchio model.
+
+    Methods:
+        setup(rmodel: pin.Model) -> None:
+            Initializes the robot model and its associated data structure for later computations.
+
+        generate(
+            initial_state: TrajectoryPoint,
+            reference_trajectory: list[TrajectoryPoint],
+        ) -> tuple[
+            npt.NDArray[np.float64],
+            list[npt.NDArray[np.float64]],
+            list[npt.NDArray[np.float64]],
+        ]:
+            Generates the initial state, reference state vectors, and control inputs for warmstart.
+
+            Parameters:
+                initial_state (TrajectoryPoint): The starting state of the robot, containing joint configuration
+                    and velocity information.
+                reference_trajectory (list[TrajectoryPoint]): A sequence of desired trajectory points, each containing
+                    joint configuration, velocity, and acceleration.
+
+            Returns:
+                tuple: A tuple containing:
+                    - x0 (npt.NDArray[np.float64]): The initial state vector `[q, v]` where `q` is the joint configuration
+                    and `v` is the joint velocity.
+                    - x_init (list[npt.NDArray[np.float64]]): A list of state vectors `[q, v]` constructed from the
+                    reference trajectory.
+                    - u_init (list[npt.NDArray[np.float64]]): A list of control inputs computed using inverse dynamics
+                    (RNEA) based on the reference trajectory.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._rmodel: pin.Model | None = None
+        self._rdata: pin.Data | None = None
+
+    def setup(self, rmodel: pin.Model) -> None:
+        self._rmodel = rmodel
+        self._rdata = self._rmodel.createData()
+
+    @override
+    def generate(
+        self,
+        initial_state: TrajectoryPoint,
+        reference_trajectory: list[TrajectoryPoint],
+    ) -> tuple[
+        npt.NDArray[np.float64],
+        list[npt.NDArray[np.float64]],
+        list[npt.NDArray[np.float64]],
+    ]:
+        """
+        Generate initial values for a warm-start of the optimization problem.
+        The state vector is `[q, v]`, where:
+            - `q` is the robot's joint configuration.
+            - `v` is the robot's joint velocity.
+        - `init_xs`: A list of state vectors `[q, v]` constructed from the reference trajectory.
+        - `init_us`: A list of control inputs computed using inverse dynamics (RNEA)
+            based on the reference trajectory.
+        """
+        # Ensure the robot model (_rmodel) is initialized before proceeding
+        assert (
+            self._rmodel is not None
+        ), "Robot model is missing in warmstart. please use warmstart.setup(rmodel)"
+
+        x0 = np.concatenate(
+            [initial_state.robot_configuration, initial_state.robot_velocity]
+        )
+        assert x0.shape[0] == (self._rmodel.nq + self._rmodel.nv), (
+            f"Expected x0 shape {(self._rmodel.nq + self._rmodel.nv)},"
+            f"from provided reference got {x0.shape}"
+        )
+
+        x_init = np.array(
+            [
+                np.hstack([point.robot_configuration, point.robot_velocity])
+                for point in reference_trajectory
+            ]
+        )
+        assert x_init.shape == (
+            len(reference_trajectory),
+            self._rmodel.nq + self._rmodel.nv,
+        ), (
+            f"Expected x_init shape {(len(reference_trajectory), self._rmodel.nq + self._rmodel.nv)}, "
+            f"from provided reference got {x_init.shape}"
+        )
+
+        u_init = [
+            pin.rnea(
+                self._rmodel,
+                self._rdata,
+                point.robot_configuration,
+                point.robot_velocity,
+                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

agimus_controller/tests/test_warm_start_base.py

@@ -0,0 +1,9 @@
+import unittest
+
+from agimus_controller.warm_start_base import WarmStartBase
+
+
+class TestWarmStartBase(unittest.TestCase):
+    def test_abstract_class_instantiation(self):
+        with self.assertRaises(TypeError):
+            WarmStartBase()

agimus_controller/tests/test_warm_start_reference.py

@@ -0,0 +1,63 @@
+import unittest
+
+import numpy as np
+import example_robot_data as robex
+import pinocchio as pin
+
+from agimus_controller.warm_start_reference import WarmStartReference
+from agimus_controller.trajectory import TrajectoryPoint
+
+
+class TestWarmStart(unittest.TestCase):
+    def test_initialization(self):
+        ws = WarmStartReference()
+        self.assertEqual(ws._previous_solution, [])
+
+    def test_generate(self):
+        ws = WarmStartReference()
+        num_points = 10
+        robot = robex.load("ur5")
+        rmodel = robot.model
+        rdata = robot.data
+        ws.setup(rmodel=rmodel)
+
+        initial_q = np.random.randn(rmodel.nq)
+        initial_v = np.random.randn(rmodel.nv)
+        initial_state = TrajectoryPoint(
+            robot_configuration=initial_q, robot_velocity=initial_v
+        )
+
+        random_qs = np.random.randn(num_points, rmodel.nq)
+        random_vs = np.random.randn(num_points, rmodel.nv)
+        random_acs = np.random.randn(num_points, rmodel.nv)
+        reference_trajectory = [
+            TrajectoryPoint(
+                robot_configuration=q, robot_velocity=v, robot_acceleration=a
+            )
+            for q, v, a in zip(random_qs, random_vs, random_acs)
+        ]
+
+        # Create the expected stacked array
+        expected_x0 = np.concatenate([initial_q, initial_v])
+        expected_x_init = np.hstack((random_qs, random_vs))
+        expected_u_init = np.array(
+            [
+                pin.rnea(rmodel, rdata, q, v, a)
+                for q, v, a in zip(random_qs, random_vs, random_acs)
+            ]
+        )
+
+        # Act
+        x0, x_init, u_init = ws.generate(initial_state, reference_trajectory)
+        x_init = np.array(x_init)
+        u_init = np.array(u_init)
+
+        # Assert
+        # Check shapes
+        self.assertEqual(x_init.shape, expected_x_init.shape)
+        self.assertEqual(u_init.shape, expected_u_init.shape)
+
+        # Check values (assuming `generate` would use these random inputs)
+        np.testing.assert_array_equal(x0, expected_x0)
+        np.testing.assert_array_equal(x_init, expected_x_init)
+        np.testing.assert_array_equal(u_init, expected_u_init)