6-DOF drone simulation with PID controller

# 발췌: DroneDynamics.update
def update(self, total_thrust, torques, dt):
  # 회전 운동: 토크로부터 각가속도를 계산
  ang_accel = self.invI @ (torques - np.cross(self.angular_velocity, self.I @ self.angular_velocity))
  self.angular_velocity += ang_accel * dt

  # 병진 운동: 모든 힘을 합산하여 가속도를 계산
  R = self.orientation.as_matrix()
  thrust_vector = R @ np.array([0, 0, total_thrust])

  # 중력, 공기 저항 계산
  net_force = thrust_vector + drag_force + gravity_force
  acceleration = net_force / self.mass
  self.velocity += acceleration * dt
  self.position += self.velocity * dt

gravity_force = np.array([0, 0, -self.mass * self.g])

# R은 self.orientation (쿼터니언) 으로부터 계산된 회전 행렬
R = self.orientation.as_matrix()
thrust_vector = R @ np.array([0, 0, total_thrust])

drag_force = -self.drag_coeff * self.velocity * np.linalg.norm(self.velocity)

# 발췌: DroneDynamics.update
# 각가속도 계산
ang_accel = self.invI @ (torques - np.cross(self.angular_velocity, self.I @ self.angular_velocity))
self.angular_velocity += ang_accel * dt

# 각속도를 이용해 쿼터니언 자세 업데이트
q_dot = self.orientation * Rotation.from_rotvec(self.angular_velocity * dt)
self.orientation = q_dot

# 발췌: generate_optimal_trajectory
# p: 최적화할 다항식 계수 벡터
p = cp.Variable(n_segments * n_coeffs)

# Objective: p^T * Q * p 형태의 이차식을 최소화
objective = cp.Minimize(cp.quad_form(p, Q_total + 1e-6 * np.eye(n_segments * n_coeffs)))

# Constraints: Ap = b 형태의 등식 제약조건
constraints = [A @ p == b[i, :]]

# 문제 정의 및 풀이
prob = cp.Problem(objective, constraints)
prob.solve(solver=cp.OSQP)

# 발췌: GeometricController.update
# 1. 목표 힘 벡터(f_des) 계산
e_pos = self.drone.position - pos_target
e_vel = self.drone.velocity - vel_target
f_des = -self.kp_pos @ e_pos - self.kd_pos @ e_vel + ...

# 2. 목표 추력(thrust_cmd) 및 목표 자세(R_des) 계산
thrust_cmd = np.dot(f_des, b3)
# ...
R_des = np.vstack([b1_des, b2_des, b3_des]).T

# 3. 자세 오차(e_R)를 이용해 최종 토크(torques_cmd) 계산
e_R_matrix = 0.5 * (R_des.T @ R - R.T @ R_des)
e_R = np.array([e_R_matrix[2,1], e_R_matrix[0,2], e_R_matrix[1,0]])
torques_cmd = -self.kp_att @ e_R - self.kd_att @ e_omega + ...

# 비용 함수: 주어진 gains로 시뮬레이션을 돌리고 성능 점수를 반환
def objective_function(gains):
  # ... 전체 시뮬레이션 실행 ...
  cost_pos = np.mean(error_pos**2)
  cost_rot = np.mean(np.array(angular_velocities)**2)
  total_cost = cost_pos + weight_rot * cost_rot
  return total_cost

# 최적화 실행: objective_function의 점수를 최소화하는 gains를 찾음
result = minimize(objective_function, initial_gains, method='Nelder-Mead', ...)
best_gains = result.x

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation
from scipy.spatial.transform import Rotation
from scipy.optimize import minimize
from IPython.display import HTML
import cvxpy as cp
import math
from typing import List, Tuple

# --- 시뮬레이션 설정 ---
# 이륙 지점 및 방문할 경유지 목록
START_POINT: np.ndarray = np.array([0., 0., 2.])
WAYPOINTS_TO_VISIT: List[np.ndarray] = [
    np.array([8., -8., 8.]),
    np.array([8., 8., 4.]),
    np.array([-8., 8., 8.]),
    np.array([-8., -8., 5.]),
    np.array([0., 0., 15.])
]

# 비행 및 시뮬레이션 파라미터
MAX_SPEED: float = 3.0
SIMULATION_TIMESTEP: float = 0.01

# 최적화 및 제어 게인 설정
INITIAL_GAINS: np.ndarray = np.array([6.0, 8.0, 4.0, 5.0, 20.0, 10.0, 2.0, 1.0])
GAIN_BOUNDS: List[Tuple[float, float]] = [(0.1, 50.0)] * 8
OPTIMIZER_MAX_ITERATIONS: int = 100

# --- 클래스 정의 ---

class DroneDynamics:
    """드론의 물리적 역학을 시뮬레이션하는 클래스."""
    def __init__(self, start_pos: np.ndarray = np.zeros(3)):
        self.mass = 1.0
        self.g = 9.81
        self.I = np.diag([0.01, 0.01, 0.02])
        self.invI = np.linalg.inv(self.I)
        self.drag_coeff = 0.1
        self.position = start_pos.copy()
        self.velocity = np.zeros(3)
        self.orientation = Rotation.from_quat([0, 0, 0, 1])
        self.angular_velocity = np.zeros(3)

    def update(self, total_thrust: float, torques: np.ndarray, dt: float) -> None:
        """한 타임스텝 동안 드론의 상태를 업데이트합니다."""
        ang_accel = self.invI @ (torques - np.cross(self.angular_velocity, self.I @ self.angular_velocity))
        self.angular_velocity += ang_accel * dt
        q_dot = self.orientation * Rotation.from_rotvec(self.angular_velocity * dt)
        self.orientation = q_dot

        R = self.orientation.as_matrix()
        thrust_vector = R @ np.array([0, 0, total_thrust])
        drag_force = -self.drag_coeff * self.velocity * np.linalg.norm(self.velocity)
        gravity_force = np.array([0, 0, -self.mass * self.g])
        net_force = thrust_vector + drag_force + gravity_force

        acceleration = net_force / self.mass
        self.velocity += acceleration * dt
        self.position += self.velocity * dt

        if self.position[2] < 0:
            self.position[2] = 0
            self.velocity[2] = 0

class Polynomial:
    """궤적의 한 조각을 나타내는 1차원 다항식 클래스."""
    def __init__(self, coeffs: np.ndarray):
        self.coeffs = coeffs

    def get_state(self, t: float) -> Tuple[float, float, float]:
        pos = np.polyval(self.coeffs, t)
        vel = np.polyval(np.polyder(self.coeffs, 1), t)
        acc = np.polyval(np.polyder(self.coeffs, 2), t)
        return pos, vel, acc

class PiecewisePolynomialTrajectory:
    """여러 다항식 조각으로 구성된 전체 비행 궤적 클래스."""
    def __init__(self, polys_xyz: List[List[Polynomial]], durations: List[float]):
        self.polys_xyz = polys_xyz
        self.durations = durations
        self.cumulative_durations = np.cumsum([0] + durations)

    def get_state(self, t: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """특정 시간 t에서의 궤적 상태(위치, 속도, 가속도)를 반환합니다."""
        if t >= self.cumulative_durations[-1]:
            t = self.cumulative_durations[-1] - 1e-6

        segment_index = np.searchsorted(self.cumulative_durations, t, side='right') - 1
        time_in_segment = t - self.cumulative_durations[segment_index]

        pos = np.array([p_list[segment_index].get_state(time_in_segment)[0] for p_list in self.polys_xyz])
        vel = np.array([p_list[segment_index].get_state(time_in_segment)[1] for p_list in self.polys_xyz])
        acc = np.array([p_list[segment_index].get_state(time_in_segment)[2] for p_list in self.polys_xyz])
        return pos, vel, acc

class GeometricController:
    """기하학적 제어기(Geometric Controller) 클래스."""
    def __init__(self, drone: DroneDynamics, gains: np.ndarray):
        self.drone = drone
        self.kp_pos = np.diag([gains[0], gains[0], gains[1]])
        self.kd_pos = np.diag([gains[2], gains[2], gains[3]])
        self.kp_att = np.diag([gains[4], gains[4], gains[5]])
        self.kd_att = np.diag([gains[6], gains[6], gains[7]])

    def update(self, trajectory_target: Tuple[np.ndarray, np.ndarray, np.ndarray], yaw_target: float, dt: float) -> Tuple[float, np.ndarray]:
        """목표 궤적을 따라가기 위한 추력과 토크를 계산합니다."""
        pos_target, vel_target, acc_target = trajectory_target
        e_pos = self.drone.position - pos_target
        e_vel = self.drone.velocity - vel_target

        f_des = (-self.kp_pos @ e_pos - self.kd_pos @ e_vel +
                 self.drone.mass * self.drone.g * np.array([0, 0, 1]) +
                 self.drone.mass * acc_target)

        R = self.drone.orientation.as_matrix()
        b3 = R[:, 2]
        thrust_cmd = np.dot(f_des, b3)

        b3_des = f_des / np.linalg.norm(f_des) if np.linalg.norm(f_des) > 1e-6 else np.array([0, 0, 1])
        a_yaw = np.array([np.cos(yaw_target), np.sin(yaw_target), 0])
        b2_des_numerator = np.cross(b3_des, a_yaw)
        b2_des = b2_des_numerator / np.linalg.norm(b2_des_numerator) if np.linalg.norm(b2_des_numerator) > 1e-6 else np.cross(b3_des, a_yaw)

        b1_des = np.cross(b2_des, b3_des)
        R_des = np.vstack([b1_des, b2_des, b3_des]).T

        e_R_matrix = 0.5 * (R_des.T @ R - R.T @ R_des)
        e_R = np.array([e_R_matrix[2, 1], e_R_matrix[0, 2], e_R_matrix[1, 0]])
        e_omega = self.drone.angular_velocity

        torques_cmd = -self.kp_att @ e_R - self.kd_att @ e_omega + np.cross(self.drone.angular_velocity, self.drone.I @ self.drone.angular_velocity)
        return thrust_cmd, torques_cmd

# --- 핵심 기능 함수 ---

def generate_optimal_trajectory(waypoints: List[np.ndarray], durations: List[float], order: int = 7) -> PiecewisePolynomialTrajectory:
    """경유지들을 통과하는 최적의(Minimum Snap) 부드러운 궤적을 생성합니다."""
    n_segments = len(waypoints) - 1
    n_coeffs = order + 1

    # QP 문제의 목적 함수 (p'Qp)를 위한 Q 행렬 구성
    Q_total = np.zeros((n_segments * n_coeffs, n_segments * n_coeffs))
    for i in range(n_segments):
        Q_i = np.zeros((n_coeffs, n_coeffs))
        T = durations[i]
        for r in range(4, n_coeffs):
            for c in range(4, n_coeffs):
                k1, k2 = order - r, order - c
                if k1 >= 4 and k2 >= 4:
                    d = (math.factorial(k1) / math.factorial(k1-4)) * (math.factorial(k2) / math.factorial(k2-4))
                    power = k1 + k2 - 7
                    if power > 0: Q_i[r, c] = 2 * d / power * (T**power)
        Q_total[i*n_coeffs:(i+1)*n_coeffs, i*n_coeffs:(i+1)*n_coeffs] = Q_i

    # 제약 조건 (Ap = b)을 위한 A, b 행렬 구성
    def get_der_coeffs(k, t, order):
        coeffs = np.zeros(order + 1)
        for i in range(order + 1):
            p = order - i
            if p >= k: coeffs[i] = math.factorial(p)/math.factorial(p-k) * (t**(p-k))
        return coeffs

    n_derivatives = 4
    n_constraints = 2 * n_derivatives + (n_segments - 1) * n_derivatives
    A = np.zeros((n_constraints, n_segments * n_coeffs))
    b = np.zeros((3, n_constraints))
    row = 0

    A[row, :n_coeffs] = get_der_coeffs(0, 0, order)
    for dim in range(3): b[dim, row] = waypoints[0][dim]
    row += 1
    for k in range(1, n_derivatives): A[row, :n_coeffs] = get_der_coeffs(k, 0, order); row += 1

    A[row, (n_segments-1)*n_coeffs:] = get_der_coeffs(0, durations[-1], order)
    for dim in range(3): b[dim, row] = waypoints[-1][dim]
    row += 1
    for k in range(1, n_derivatives): A[row, (n_segments-1)*n_coeffs:] = get_der_coeffs(k, durations[-1], order); row += 1

    for i in range(n_segments - 1):
        T_i = durations[i]
        A[row, i*n_coeffs:(i+1)*n_coeffs] = get_der_coeffs(0, T_i, order)
        for dim in range(3): b[dim, row] = waypoints[i+1][dim]
        row += 1
        for k in range(1, n_derivatives):
            A[row, i*n_coeffs:(i+1)*n_coeffs] = get_der_coeffs(k, T_i, order)
            A[row, (i+1)*n_coeffs:(i+2)*n_coeffs] = -get_der_coeffs(k, 0, order)
            row += 1

    # 각 축(x, y, z)에 대해 QP 문제 풀이
    polys_xyz = [[], [], []]
    for i in range(3):
        p = cp.Variable(n_segments * n_coeffs)
        objective = cp.Minimize(cp.quad_form(p, Q_total + 1e-6 * np.eye(n_segments * n_coeffs)))
        constraints = [A @ p == b[i, :]]
        prob = cp.Problem(objective, constraints)
        prob.solve(solver=cp.OSQP, verbose=False, max_iter=50000)
        if prob.status not in ['optimal', 'optimal_inaccurate']:
            raise ValueError(f"최적화 실패 (축: {i}, 상태: {prob.status})")
        coeffs_all = p.value
        for j in range(n_segments):
            polys_xyz[i].append(Polynomial(coeffs_all[j*n_coeffs:(j+1)*n_coeffs]))
    return PiecewisePolynomialTrajectory(polys_xyz, durations)

def plan_tsp_mission(waypoints_to_visit: List[np.ndarray], start_end_point: np.ndarray, max_speed: float) -> Tuple[List[np.ndarray], List[float]]:
    """최근접 이웃 알고리즘으로 여러 경유지의 최적 방문 순서와 비행 시간을 계획합니다."""
    unvisited = list(waypoints_to_visit)
    ordered_path = [start_end_point]
    current_node = start_end_point
    while unvisited:
        distances = [np.linalg.norm(current_node - node) for node in unvisited]
        nearest_index = np.argmin(distances)
        current_node = unvisited.pop(nearest_index)
        ordered_path.append(current_node)
    ordered_path.append(start_end_point)

    durations = [max(np.linalg.norm(p2 - p1) / max_speed, 3.0) for p1, p2 in zip(ordered_path[:-1], ordered_path[1:])]
    return ordered_path, durations

def objective_function(gains: np.ndarray) -> float:
    """제어 게인 튜닝을 위한 목적 함수. 시뮬레이션 오차를 비용으로 반환합니다."""
    drone = DroneDynamics(start_pos=START_POINT)
    controller = GeometricController(drone, gains)
    positions, angular_velocities = [], []
    last_yaw_target = 0.0

    for i in range(timesteps):
        t = i * SIMULATION_TIMESTEP
        pos_target, vel_target, acc_target = trajectory.get_state(t)

        if np.linalg.norm(vel_target[:2]) > 0.1:
            current_yaw_target = np.arctan2(vel_target[1], vel_target[0])
            diff = current_yaw_target - last_yaw_target
            if diff > np.pi: last_yaw_target += 2 * np.pi
            elif diff < -np.pi: last_yaw_target -= 2 * np.pi
            target_yaw = 0.9 * last_yaw_target + 0.1 * current_yaw_target
            last_yaw_target = target_yaw
        else:
            target_yaw = last_yaw_target

        thrust, torques = controller.update((pos_target, vel_target, acc_target), target_yaw, SIMULATION_TIMESTEP)
        drone.update(thrust, torques, SIMULATION_TIMESTEP)
        positions.append(drone.position.copy())
        angular_velocities.append(drone.angular_velocity.copy())

    cost_pos = np.mean((np.array(positions) - ref_path)**2)
    cost_rot = np.mean(np.array(angular_velocities)**2)
    total_cost = cost_pos + 0.01 * cost_rot

    print(f"Cost: {total_cost:.6f} | Gains: {[f'{g:.2f}' for g in gains]}")
    return total_cost

# --- 메인 실행 로직 ---

# 1. 임무 계획 및 궤적 생성
print("1. 임무 계획 및 궤적 생성 중...")
waypoints, durations = plan_tsp_mission(WAYPOINTS_TO_VISIT, START_POINT, MAX_SPEED)
trajectory = generate_optimal_trajectory(waypoints, durations)
sim_time = sum(durations)
timesteps = int(sim_time / SIMULATION_TIMESTEP)
ref_path = np.array([trajectory.get_state(i * SIMULATION_TIMESTEP)[0] for i in range(timesteps)])

# 2. 제어 게인 자동 튜닝
print("\n2. 제어 게인 자동 튜닝 시작...")
result = minimize(objective_function, INITIAL_GAINS, method='Nelder-Mead', bounds=GAIN_BOUNDS,
                  options={'disp': True, 'maxiter': OPTIMIZER_MAX_ITERATIONS})
best_gains = result.x
print(f"\n튜닝 완료! 최적 게인: {[f'{g:.3f}' for g in best_gains]}")

# 3. 최적 게인으로 최종 시뮬레이션 실행
print("\n3. 최적 게인으로 최종 시뮬레이션 실행 중...")
drone = DroneDynamics(start_pos=START_POINT)
controller = GeometricController(drone, gains=best_gains)
history = {'time': [], 'pos': [], 'rot': []}
last_yaw_target = 0.0
for i in range(timesteps):
    t = i * SIMULATION_TIMESTEP
    pos_target, vel_target, acc_target = trajectory.get_state(t)
    if np.linalg.norm(vel_target[:2]) > 0.1:
        current_yaw_target = np.arctan2(vel_target[1], vel_target[0])
        diff = current_yaw_target - last_yaw_target
        if diff > np.pi: last_yaw_target += 2 * np.pi
        elif diff < -np.pi: last_yaw_target -= 2 * np.pi
        target_yaw = 0.9 * last_yaw_target + 0.1 * current_yaw_target; last_yaw_target = target_yaw
    else: target_yaw = last_yaw_target
    thrust, torques = controller.update((pos_target, vel_target, acc_target), target_yaw, SIMULATION_TIMESTEP)
    drone.update(thrust, torques, SIMULATION_TIMESTEP)
    history['time'].append(t)
    history['pos'].append(drone.position.copy())
    history['rot'].append(drone.orientation.as_matrix())

# 4. 애니메이션 생성
print("\n4. 최종 애니메이션 생성 중...")
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(111, projection='3d')
body_size = 0.3
drone_body = np.array([[body_size, 0, 0], [-body_size, 0, 0], [0, body_size, 0], [0, -body_size, 0]]).T

all_points = np.array(waypoints)
padding = 3.0
ax_limits = [
    (all_points[:,0].min() - padding, all_points[:,0].max() + padding),
    (all_points[:,1].min() - padding, all_points[:,1].max() + padding),
    (0, all_points[:,2].max() + padding)
]

def update_animation(i: int):
    frame_idx = min(i * (timesteps // 500), len(history['pos']) - 1)
    ax.cla()
    pos_now, R_now = history['pos'][frame_idx], history['rot'][frame_idx]

    drone_body_world = (R_now @ drone_body).T + pos_now
    ax.plot(drone_body_world[:2,0], drone_body_world[:2,1], drone_body_world[:2,2], 'k-', lw=3)
    ax.plot(drone_body_world[2:,0], drone_body_world[2:,1], drone_body_world[2:,2], 'k-', lw=3)

    front_marker = (R_now @ np.array([body_size, 0, 0])).T + pos_now
    ax.scatter(front_marker[0], front_marker[1], front_marker[2], color='red', s=40, label='Front')

    ax.plot(np.array(history['pos'])[:frame_idx,0], np.array(history['pos'])[:frame_idx,1], np.array(history['pos'])[:frame_idx,2], 'b-', lw=2.5, label='실제 궤적')
    ax.plot(ref_path[:,0], ref_path[:,1], ref_path[:,2], 'g--', lw=2, label='목표 궤적')

    for j, wp in enumerate(waypoints):
        ax.scatter(wp[0], wp[1], wp[2], s=150, marker='s' if j in [0, len(waypoints)-1] else 'o', label=f'경유지 {j}', depthshade=False)

    ax.set(xlabel='X (m)', ylabel='Y (m)', zlabel='Z (m)',
           title=f"자동 튜닝된 드론 임무 | 시간: {history['time'][frame_idx]:.2f}s",
           xlim=ax_limits[0], ylim=ax_limits[1], zlim=ax_limits[2])
    ax.legend()
    ax.view_init(elev=30, azim=-110)

ani = FuncAnimation(fig, update_animation, frames=500, interval=40)
plt.close()
display(HTML(ani.to_html5_video()))
print("\n애니메이션 준비 완료.")