diff --git a/agimus_controller/agimus_controller/trajectory.py b/agimus_controller/agimus_controller/trajectory.py index 10332ad3..a7de02a9 100644 --- a/agimus_controller/agimus_controller/trajectory.py +++ b/agimus_controller/agimus_controller/trajectory.py @@ -1,6 +1,7 @@ from collections import deque from dataclasses import dataclass import numpy as np +import numpy.typing as npt from pinocchio import SE3, Force @@ -8,25 +9,25 @@ 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 diff --git a/agimus_controller/agimus_controller/warm_start_base.py b/agimus_controller/agimus_controller/warm_start_base.py index 2950517a..6ff529ec 100644 --- a/agimus_controller/agimus_controller/warm_start_base.py +++ b/agimus_controller/agimus_controller/warm_start_base.py @@ -1,24 +1,92 @@ +"""WarmStartBase Module. + +This module defines the WarmStartBase class, a template for generating warm-starts +for optimal control problems. It includes methods to initialize state and control +trajectories based on an initial robot state and a reference trajectory. + +Example: + Subclass the WarmStartBase class and implement the `generate` method to create + a warm-start for a specific optimization problem. +""" + from abc import ABC, abstractmethod + import numpy as np +import numpy.typing as npt from agimus_controller.trajectory import TrajectoryPoint class WarmStartBase(ABC): + """Base class for generating warm-starts for optimal control problems. + + This class provides a template for generating initial values for state and + control trajectories based on the initial robot state and a reference trajectory. + """ + def __init__(self) -> None: + """Initialize the WarmStartBase class.""" super().__init__() - self._previous_solution: list[TrajectoryPoint] = list() + # Stores the previous solution of the optimization problem. + self._previous_solution: list[TrajectoryPoint] = [] @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: + - npt.NDArray[np.float64]: The initial state vector. + - list[npt.NDArray[np.float64]]: List of state vectors + for each point in the reference trajectory. + - list[npt.NDArray[np.float64]]: List of control inputs. + """ + ... + + @abstractmethod + def setup(self, *args, **kwargs) -> None: + """Sets up the variables needed for the warmstart computation. Allows to pass additional variables after the class is initialised, that are only know at the runtime. + + Args: + *args: Variable length argument list. + **kwargs: Arbitrary keyword arguments. + + Example: + >>> class MyPinocchioWarmstart(WarmStartBase): + ... def setup(self, rmodel: pin.Model) -> None: + ... self._rmodel = rmodel + ... self._rdata = self._rmodel.createData() + ... self._nx = self._rmodel.nq + self._rmodel.nv + >>> warmstart = MyPinocchioWarmstart() + >>> # ... + >>> rmodel: pin.Model() = await_urdf_model() + >>> warmstart.setup(rmodel) + """ ... def update_previous_solution( self, previous_solution: list[TrajectoryPoint] ) -> None: - """Stores internally the previous solution of the OCP""" + """Update the stored previous solution. + + Stores the solution from a previous optimization cycle to be used as a reference + or initialization for warm-start generation. + Args: + previous_solution (list[TrajectoryPoint]): The solution of the optimization problem. + """ self._previous_solution = previous_solution diff --git a/agimus_controller/agimus_controller/warm_start_reference.py b/agimus_controller/agimus_controller/warm_start_reference.py new file mode 100644 index 00000000..77009e4e --- /dev/null +++ b/agimus_controller/agimus_controller/warm_start_reference.py @@ -0,0 +1,91 @@ +import numpy as np +import numpy.typing as npt + +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. + """ + + def __init__(self) -> None: + super().__init__() + # The robot's Pinocchio model, used for forward dynamics computations. + self._rmodel: pin.Model | None = None + # Data structure associated with the Pinocchio model. + self._rdata: pin.Data | None = None + # Size of the state vector + self._nx: int = 0 + + def setup(self, rmodel: pin.Model) -> None: + self._rmodel = rmodel + self._rdata = self._rmodel.createData() + self._nx = self._rmodel.nq + self._rmodel.nv + + 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._nx), ( + f"Expected x0 shape {(self._nx)}," 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._nx, + ), ( + f"Expected x_init shape {(len(reference_trajectory), self._nx)}, " + 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 diff --git a/agimus_controller/tests/__init__.py b/agimus_controller/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/agimus_controller/tests/test_warm_start_base.py b/agimus_controller/tests/test_warm_start_base.py new file mode 100644 index 00000000..28774952 --- /dev/null +++ b/agimus_controller/tests/test_warm_start_base.py @@ -0,0 +1,13 @@ +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() + + +if __name__ == "__main__": + unittest.main() diff --git a/agimus_controller/tests/test_warm_start_reference.py b/agimus_controller/tests/test_warm_start_reference.py new file mode 100644 index 00000000..45f1689f --- /dev/null +++ b/agimus_controller/tests/test_warm_start_reference.py @@ -0,0 +1,67 @@ +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) + + +if __name__ == "__main__": + unittest.main()