src/sqp_optimizer/__init__.py - develop - oop_nust_misis/NumpyBasics-Plesik

src/sqp_optimizer/__init__.py
439 строк16 KB
vivianlo
src done
17 фев 2026, 14:37
17 фев 2026, 14:375939a92
from typing import Dict, Optional, Tuple

import numpy as np
from matplotlib import pyplot as plt

## TODO: Добавьте необходимые импорты для работы с данными и алгоритмами
from scipy.optimize import minimize


class SQPTrajectoryOptimizer:
    """
    Sequential Quadratic Programming оптимизатор для траекторий мобильного робота
    с учетом ограничений на состояние и управление.

    Студенты должны реализовать все методы, отмеченные комментарием TODO.
    """

    def __init__(self, dt: float = 0.1, N: int = 50):
        """
        Инициализация оптимизатора

        Args:
            dt: Шаг дискретизации времени
            N: Количество шагов предсказания
        """
        self.dt = dt
        self.N = N

        # Размерности
        self.nx = 5  # состояние: [x, y, yaw, v, w]
        self.nu = 2  # управление: [a, α] (линейное и угловое ускорение)

        # Параметры по умолчанию
        self.Q = np.diag([10, 10, 5, 1, 1])  # веса состояния
        self.R = np.diag([0.1, 0.1])  # веса управления
        self.Q_terminal = np.diag([100, 100, 50, 10, 10])  # терминальные веса

        # Ограничения по умолчанию
        self.x_min = np.array([-10, -10, -np.pi, -2, -1])
        self.x_max = np.array([10, 10, np.pi, 2, 1])
        self.u_min = np.array([-1, -0.5])
        self.u_max = np.array([1, 0.5])

        self.solution = None
        self.optimization_result = None

    # --- служебные методы, чтобы удобно упаковывать/распаковывать переменные ---
    def _unpack_z(self, z: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Разложить общий вектор z на:
        x_traj shape (N+1, nx)
        u_traj shape (N, nu)

        Формат z:
        [x0, u0, x1, u1, ..., x_{N-1}, u_{N-1}, x_N]
        """
        z = np.asarray(z).ravel()
        x_traj = np.zeros((self.N + 1, self.nx))
        u_traj = np.zeros((self.N, self.nu))

        idx = 0
        for k in range(self.N):
            x_traj[k] = z[idx : idx + self.nx]
            idx += self.nx
            u_traj[k] = z[idx : idx + self.nu]
            idx += self.nu

        x_traj[self.N] = z[idx : idx + self.nx]
        return x_traj, u_traj

    def _pack_z(self, x_traj: np.ndarray, u_traj: np.ndarray) -> np.ndarray:
        """
        Собрать z обратно из x_traj и u_traj в том же порядке:
        [x0, u0, x1, u1, ..., x_N]
        """
        parts = []
        for k in range(self.N):
            parts.append(np.asarray(x_traj[k]).ravel())
            parts.append(np.asarray(u_traj[k]).ravel())
        parts.append(np.asarray(x_traj[self.N]).ravel())
        return np.concatenate(parts)

    def extended_dynamics(self, x: np.ndarray, u: np.ndarray) -> np.ndarray:
        """
        Расширенная динамика мобильного робота с ускорениями

        TODO: Реализовать модель динамики робота

        Args:
            x: Текущее состояние [x, y, yaw, v, w]
            u: Управление [a, α] (ускорения)

        Returns:
            Новое состояние после применения динамики
        """
        # TODO: Реализовать дискретную модель динамики:
        # x_new[0] = x[0] + x[3] * cos(x[2]) * dt  # позиция x
        # x_new[1] = x[1] + x[3] * sin(x[2]) * dt  # позиция y
        # x_new[2] = x[2] + x[4] * dt              # угол yaw
        # x_new[3] = x[3] + u[0] * dt              # линейная скорость
        # x_new[4] = x[4] + u[1] * dt              # угловая скорость

        # Простыми словами: берём состояние и делаем шаг dt вперёд по формуле выше
        x = np.asarray(x).ravel()
        u = np.asarray(u).ravel()

        x_new = np.zeros_like(x)
        x_new[0] = x[0] + x[3] * np.cos(x[2]) * self.dt
        x_new[1] = x[1] + x[3] * np.sin(x[2]) * self.dt
        x_new[2] = x[2] + x[4] * self.dt
        x_new[3] = x[3] + u[0] * self.dt
        x_new[4] = x[4] + u[1] * self.dt

        # чтобы yaw не рос бесконечно, нормируем в [-pi, pi]
        x_new[2] = (x_new[2] + np.pi) % (2 * np.pi) - np.pi

        return x_new

    def cost_function(self, z: np.ndarray, x_ref: np.ndarray) -> float:
        """
        Функция стоимости для SQP

        TODO: Реализовать функцию стоимости для задачи оптимизации

        Args:
            z: Вектор оптимизации [x0, u0, x1, u1, ..., x_N]
            x_ref: Опорная траектория

        Returns:
            Значение функции стоимости
        """
        # TODO: Реализовать вычисление стоимости:
        # cost = sum_{k=0}^{N-1} [ (x_k - x_ref[k])^T @ Q @ (x_k - x_ref[k]) + u_k^T @ R @ u_k ]
        #        + (x_N - x_ref[N])^T @ Q_terminal @ (x_N - x_ref[N])

        # Идея: штрафуем (1) отклонение от опорной траектории и (2) большие управления
        x_traj, u_traj = self._unpack_z(z)

        cost = 0.0
        for k in range(self.N):
            dx = x_traj[k] - x_ref[k]
            cost += float(dx.T @ self.Q @ dx)
            uk = u_traj[k]
            cost += float(uk.T @ self.R @ uk)

        dxN = x_traj[self.N] - x_ref[self.N]
        cost += float(dxN.T @ self.Q_terminal @ dxN)

        return cost

    def dynamics_constraints(self, z: np.ndarray, x0: Optional[np.ndarray] = None) -> np.ndarray:
        """
        Ограничения динамики системы

        TODO: Реализовать ограничения, обеспечивающие выполнение динамики

        Args:
            z: Вектор оптимизации

        Returns:
            Вектор ограничений динамики (должны быть равны 0)
        """
        # TODO: Для каждого k от 0 до N-1 добавить ограничение:
        # x_{k+1} - f(x_k, u_k) = 0

        x_traj, u_traj = self._unpack_z(z)
        cons = []

        # Фиксируем начальное состояние: x_0 = x0 (это равенство тоже должно быть 0)
        if x0 is not None:
            x0 = np.asarray(x0).ravel()
            cons.append(x_traj[0] - x0)

        # Основные ограничения динамики
        for k in range(self.N):
            x_next_pred = self.extended_dynamics(x_traj[k], u_traj[k])
            cons.append(x_traj[k + 1] - x_next_pred)

        return np.concatenate(cons)

    def state_constraints(self, z: np.ndarray) -> np.ndarray:
        """
        Ограничения на состояние

        TODO: Реализовать ограничения вида x_min <= x_k <= x_max

        Args:
            z: Вектор оптимизации

        Returns:
            Вектор ограничений состояния (должны быть >= 0)
        """
        # TODO: Для каждого состояния x_k добавить два ограничения:
        # x_k - x_min >= 0
        # x_max - x_k >= 0
        x_traj, _ = self._unpack_z(z)

        g = []
        for k in range(self.N + 1):
            g.append(x_traj[k] - self.x_min)  # не ниже минимума
            g.append(self.x_max - x_traj[k])  # не выше максимума
        return np.concatenate(g)

    def control_constraints(self, z: np.ndarray) -> np.ndarray:
        """
        Ограничения на управление

        TODO: Реализовать ограничения вида u_min <= u_k <= u_max

        Args:
            z: Вектор оптимизации

        Returns:
            Вектор ограничений управления (должны быть >= 0)
        """
        # TODO: Для каждого управления u_k добавить два ограничения:
        # u_k - u_min >= 0
        # u_max - u_k >= 0

        _, u_traj = self._unpack_z(z)

        g = []
        for k in range(self.N):
            g.append(u_traj[k] - self.u_min)  # не меньше u_min
            g.append(self.u_max - u_traj[k])  # не больше u_max
        return np.concatenate(g)

    def solve(
        self,
        x0: np.ndarray,
        x_ref: Optional[np.ndarray] = None,
        method: str = "SLSQP",
        max_iter: int = 1000,
    ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
        """
        Решение задачи траекторной оптимизации

        TODO: Реализовать настройку и решение задачи оптимизации

        Args:
            x0: Начальное состояние
            x_ref: Опорная траектория (если None, используется нулевая)
            method: Метод оптимизации
            max_iter: Максимальное количество итераций

        Returns:
            x_traj, u_traj: Оптимальная траектория и управления
        """
        x0 = np.asarray(x0).ravel()

        if x_ref is None:
            x_ref = np.zeros((self.N + 1, self.nx))
        else:
            x_ref = np.asarray(x_ref)
            if x_ref.shape != (self.N + 1, self.nx):
                raise ValueError(f"x_ref должен иметь форму {(self.N + 1, self.nx)}")

        # TODO: Создать начальное предположение z0
        x_init = x_ref.copy()
        x_init[0] = x0
        u_init = np.zeros((self.N, self.nu))
        z0 = self._pack_z(x_init, u_init)

        # TODO: Определить ограничения для scipy.optimize.minimize
        constraints = [
            {"type": "eq", "fun": lambda z, x0=x0: self.dynamics_constraints(z, x0=x0)},
            {"type": "ineq", "fun": self.state_constraints},
            {"type": "ineq", "fun": self.control_constraints},
        ]

        # (Полезно) bounds: прямые границы для переменных, чтобы оптимизатор не улетал
        bounds = []
        for k in range(self.N):
            # bounds для x_k
            if k == 0:
                # x0 фиксируем точно: (x0[i], x0[i])
                for i in range(self.nx):
                    bounds.append((float(x0[i]), float(x0[i])))
            else:
                for i in range(self.nx):
                    bounds.append((float(self.x_min[i]), float(self.x_max[i])))

            # bounds для u_k
            for i in range(self.nu):
                bounds.append((float(self.u_min[i]), float(self.u_max[i])))

        # bounds для x_N
        for i in range(self.nx):
            bounds.append((float(self.x_min[i]), float(self.x_max[i])))

        # TODO: Решить задачу оптимизации
        res = minimize(
            fun=lambda z: self.cost_function(z, x_ref),
            x0=z0,
            method=method,
            constraints=constraints,
            bounds=bounds,
            options={"maxiter": max_iter, "disp": False},
        )

        self.optimization_result = res

        if not res.success:
            self.solution = None
            return None, None

        self.solution = res.x

        # TODO: Распаковать решение в x_traj и u_traj
        x_traj, u_traj = self._unpack_z(res.x)
        return x_traj, u_traj

    def analyze_solution(self, x_traj: np.ndarray, u_traj: np.ndarray) -> Dict:
        """
        Анализ оптимального решения

        TODO: Реализовать вычисление метрик качества решения

        Args:
            x_traj: Оптимальная траектория
            u_traj: Оптимальные управления

        Returns:
            Словарь с метриками решения
        """
        # TODO: Вычислить метрики:

        x_traj = np.asarray(x_traj)
        u_traj = np.asarray(u_traj)

        # Энергия управления
        energy_u = float(np.sum(u_traj**2))

        # Максимальные по модулю управления
        max_u = np.max(np.abs(u_traj), axis=0)
        max_a = float(max_u[0])
        max_alpha = float(max_u[1])

        # Ошибка в конечной точке: расстояние конечного состояния до (0,0,0,0,0)
        final_error = float(np.linalg.norm(x_traj[-1]))

        # Количество нарушений ограничений (сколько элементов вышли за пределы)
        state_viol = int(np.sum((x_traj < self.x_min) | (x_traj > self.x_max)))
        control_viol = int(np.sum((u_traj < self.u_min) | (u_traj > self.u_max)))
        total_violations = state_viol + control_viol

        # Плавность управления: насколько управления скачут между шагами
        du = np.diff(u_traj, axis=0)
        smoothness = float(np.sum(du**2))

        return {
            "energy_u": energy_u,
            "max_a": max_a,
            "max_alpha": max_alpha,
            "final_state_error_to_zero": final_error,
            "total_violations_count": total_violations,
            "control_smoothness": smoothness,
        }


def create_reference_trajectory(
    N: int, dt: float, traj_type: str = "straight"
) -> np.ndarray:
    """
    Создание опорной траектории

    Args:
        N: Количество точек
        dt: Шаг времени
        traj_type: Тип траектории ('straight', 'curve')

    Returns:
        Опорная траектория
    """
    x_ref = np.zeros((N + 1, 5))

    if traj_type == "straight":
        for i in range(N + 1):
            alpha = i / N
            x_ref[i, 0] = 5 * (1 - alpha)
            x_ref[i, 1] = 3 * (1 - alpha)

    elif traj_type == "curve":
        t = np.arange(N + 1) * dt
        x_ref[:, 0] = 5 * np.cos(0.2 * t)
        x_ref[:, 1] = 3 * np.sin(0.2 * t)

    return x_ref


# TODO: Реализовать создание опорной траектории
def main():
    N = 100
    dt = 0.1
    x_ref = create_reference_trajectory(N, dt)

    # Пример запуска оптимизатора
    opt = SQPTrajectoryOptimizer(dt=dt, N=N)

    # Начальное состояние: возьмём первую точку опорной траектории
    # [x, y, yaw, v, w] — yaw/v/w пусть будут 0
    x0 = np.array([x_ref[0, 0], x_ref[0, 1], 0.0, 0.0, 0.0])

    x_traj, u_traj = opt.solve(x0=x0, x_ref=x_ref, method="SLSQP", max_iter=500)

    # Если не получилось — просто покажем x_ref
    if x_traj is None or u_traj is None:
        plt.figure()
        plt.plot(x_ref[:, 0], x_ref[:, 1], label="x_ref")
        plt.title("Reference trajectory (optimization failed)")
        plt.axis("equal")
        plt.legend()
        plt.show()
        return

    metrics = opt.analyze_solution(x_traj, u_traj)
    print("Metrics:", metrics)

    # График траекторий
    plt.figure()
    plt.plot(x_ref[:, 0], x_ref[:, 1], label="x_ref")
    plt.plot(x_traj[:, 0], x_traj[:, 1], label="x_opt")
    plt.title("Reference vs Optimized trajectory")
    plt.axis("equal")
    plt.legend()
    plt.show()

    # График управлений
    t_u = np.arange(N) * dt
    plt.figure()
    plt.plot(t_u, u_traj[:, 0], label="a (linear accel)")
    plt.plot(t_u, u_traj[:, 1], label="alpha (angular accel)")
    plt.title("Control inputs")
    plt.legend()
    plt.show()


if __name__ == "__main__":
    main()