286 lines
10 KiB
Python
286 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import time
|
|
import math
|
|
import argparse
|
|
import numpy as np
|
|
from time import sleep, perf_counter
|
|
from typing import Tuple
|
|
from pymavlink import mavutil
|
|
import pymavlink.dialects.v20.ardupilotmega as mavlink
|
|
from pymavlink.dialects.v20.ardupilotmega import (
|
|
MAVLink_set_position_target_local_ned_message,
|
|
MAV_FRAME_BODY_OFFSET_NED,
|
|
MAV_FRAME_LOCAL_NED,
|
|
MAV_CMD_DO_SET_MODE,
|
|
MAV_CMD_NAV_TAKEOFF,
|
|
MAV_CMD_NAV_LAND,
|
|
)
|
|
|
|
HOLD_ALT = 5.0
|
|
LOWEST_ALT = 3.0
|
|
GUIDED_MODE = 4
|
|
GUIDED_NOGPS_MODE = 20
|
|
|
|
DEFAULT_WPNAV_SPEED = 150
|
|
DEFAULT_WPNAV_ACCEL = 100
|
|
DEFAULT_WPNAV_SPEED_UP = 100
|
|
DEFAULT_WPNAV_SPEED_DN = 75
|
|
DEFAULT_WPNAV_ACCEL_Z = 75
|
|
DEFAULT_LOIT_SPEED = 150
|
|
|
|
|
|
class Controller:
|
|
HOLD_ALT = HOLD_ALT
|
|
LOWEST_ALT = LOWEST_ALT
|
|
|
|
def __init__(self, connection_string: str):
|
|
print(f"[UAV] Connecting to {connection_string} ...")
|
|
self.conn = mavutil.mavlink_connection(connection_string)
|
|
self.conn.wait_heartbeat()
|
|
print("[UAV] Heartbeat from system %u component %u" %
|
|
(self.conn.target_system, self.conn.target_component))
|
|
|
|
self.conn.mav.request_data_stream_send(
|
|
self.conn.target_system,
|
|
self.conn.target_component,
|
|
mavutil.mavlink.MAV_DATA_STREAM_ALL,
|
|
50, 1)
|
|
|
|
self.current_yaw = 0.0
|
|
self.altitude = 0.0
|
|
self.position = {'x': 0.0, 'y': 0.0, 'z': 0.0}
|
|
self.armed = False
|
|
|
|
def _drain_messages(self):
|
|
while True:
|
|
msg = self.conn.recv_match(blocking=False)
|
|
if msg is None:
|
|
break
|
|
mtype = msg.get_type()
|
|
if mtype == 'ATTITUDE':
|
|
self.current_yaw = msg.yaw
|
|
elif mtype == 'LOCAL_POSITION_NED':
|
|
self.position = {'x': msg.x, 'y': msg.y, 'z': msg.z}
|
|
self.altitude = -msg.z
|
|
elif mtype == 'GLOBAL_POSITION_INT':
|
|
if self.altitude == 0:
|
|
self.altitude = msg.relative_alt / 1000.0
|
|
elif mtype == 'HEARTBEAT':
|
|
self.armed = bool(msg.base_mode & mavutil.mavlink.MAV_MODE_FLAG_SAFETY_ARMED)
|
|
elif mtype == 'STATUSTEXT':
|
|
text = msg.text if isinstance(msg.text, str) else msg.text.decode('utf-8', errors='ignore')
|
|
text = text.strip()
|
|
if text:
|
|
print(f"[SITL] {text}")
|
|
|
|
def update_state(self):
|
|
self._drain_messages()
|
|
|
|
def get_local_position(self) -> dict:
|
|
self.update_state()
|
|
return self.position.copy()
|
|
|
|
def get_altitude(self) -> float:
|
|
self.update_state()
|
|
return self.altitude
|
|
|
|
def set_param(self, name: str, value: float):
|
|
self.conn.mav.param_set_send(
|
|
self.conn.target_system,
|
|
self.conn.target_component,
|
|
name.encode('utf-8'),
|
|
value,
|
|
mavutil.mavlink.MAV_PARAM_TYPE_REAL32)
|
|
sleep(0.1)
|
|
self._drain_messages()
|
|
|
|
def configure_ekf_fast_converge(self):
|
|
"""Tune EKF parameters for faster initial convergence in SITL."""
|
|
params = {
|
|
'EK3_CHECK_SCALE': 200, # Relax EKF health checks (default 100)
|
|
'EK3_GPS_CHECK': 0, # Disable GPS preflight checks in SITL
|
|
'EK3_POSNE_M_NSE': 0.5, # Trust GPS position more
|
|
'EK3_VELD_M_NSE': 0.2, # Trust GPS vertical vel more (default 0.4)
|
|
'EK3_VELNE_M_NSE': 0.1, # Trust GPS horizontal vel more (default 0.3)
|
|
'INS_GYRO_CAL': 0, # Skip gyro cal on boot (SITL doesn't need it)
|
|
'ARMING_CHECK': 0, # Disable ALL preflight checks (SITL only)
|
|
}
|
|
for name, value in params.items():
|
|
self.set_param(name, value)
|
|
print("[UAV] EKF fast-converge params applied")
|
|
|
|
def set_mode_guided(self):
|
|
self.conn.mav.command_long_send(
|
|
self.conn.target_system,
|
|
self.conn.target_component,
|
|
MAV_CMD_DO_SET_MODE,
|
|
0, 89, GUIDED_MODE, 0, 0, 0, 0, 0)
|
|
sleep(1)
|
|
self._drain_messages()
|
|
print("[UAV] Mode -> GUIDED")
|
|
|
|
def set_mode_guided_nogps(self):
|
|
self.conn.mav.command_long_send(
|
|
self.conn.target_system,
|
|
self.conn.target_component,
|
|
MAV_CMD_DO_SET_MODE,
|
|
0, 89, GUIDED_NOGPS_MODE, 0, 0, 0, 0, 0)
|
|
sleep(1)
|
|
self._drain_messages()
|
|
print("[UAV] Mode -> GUIDED_NOGPS")
|
|
|
|
def wait_for_ekf(self, timeout: float = 120.0) -> bool:
|
|
"""Wait until EKF has actual GPS-fused position estimate."""
|
|
print("[UAV] Waiting for EKF position estimate ...")
|
|
t0 = time.time()
|
|
while time.time() - t0 < timeout:
|
|
self._drain_messages()
|
|
msg = self.conn.recv_match(type='EKF_STATUS_REPORT', blocking=True, timeout=2)
|
|
if msg is not None:
|
|
# bit 0x08 = pos_horiz_abs (actual GPS-fused position)
|
|
# bit 0x20 = pred_pos_horiz_abs (predicted, not enough for arming)
|
|
pos_horiz_abs = bool(msg.flags & 0x08)
|
|
if pos_horiz_abs:
|
|
print(f"\n[UAV] EKF has position estimate (flags=0x{msg.flags:04x})")
|
|
return True
|
|
elapsed = int(time.time() - t0)
|
|
print(f"\r[UAV] Waiting for EKF ... {elapsed}s ", end='', flush=True)
|
|
print("\n[UAV] EKF wait timed out")
|
|
return False
|
|
|
|
def arm(self, retries: int = 15):
|
|
self.set_mode_guided()
|
|
|
|
# Wait for EKF before burning through arm attempts
|
|
if not self.wait_for_ekf(timeout=120.0):
|
|
print("[UAV] Aborting arm: EKF never became healthy")
|
|
return False
|
|
|
|
for i in range(retries):
|
|
print(f"[UAV] Arm attempt {i+1}/{retries}...")
|
|
self.conn.mav.command_long_send(
|
|
self.conn.target_system,
|
|
self.conn.target_component,
|
|
mavutil.mavlink.MAV_CMD_COMPONENT_ARM_DISARM,
|
|
0, 1, 21196, 0, 0, 0, 0, 0)
|
|
|
|
armed = False
|
|
for _ in range(20):
|
|
self._drain_messages()
|
|
if self.conn.motors_armed():
|
|
armed = True
|
|
break
|
|
time.sleep(0.2)
|
|
if armed:
|
|
self.armed = True
|
|
print("[UAV] Armed")
|
|
return True
|
|
print(f"[UAV] Arm attempt {i+1} failed, retrying in 2s...")
|
|
time.sleep(2)
|
|
return False
|
|
|
|
def disarm(self):
|
|
self.conn.arducopter_disarm()
|
|
print("[UAV] Disarmed")
|
|
|
|
def takeoff(self, altitude: float = HOLD_ALT):
|
|
self.conn.mav.command_long_send(
|
|
self.conn.target_system,
|
|
self.conn.target_component,
|
|
MAV_CMD_NAV_TAKEOFF,
|
|
0, 0, 0, 0, 0, 0, 0, altitude)
|
|
print(f"[UAV] Takeoff -> {altitude}m")
|
|
|
|
def land(self):
|
|
self.conn.mav.command_long_send(
|
|
self.conn.target_system,
|
|
self.conn.target_component,
|
|
MAV_CMD_DO_SET_MODE,
|
|
0, 89, 9, 0, 0, 0, 0, 0)
|
|
print("[UAV] Landing (LAND mode)")
|
|
|
|
def move_pos_rel(self, x: float, y: float, z: float):
|
|
move_msg = MAVLink_set_position_target_local_ned_message(
|
|
int(perf_counter() * 1000),
|
|
self.conn.target_system,
|
|
self.conn.target_component,
|
|
MAV_FRAME_BODY_OFFSET_NED,
|
|
3576, x, y, z,
|
|
0, 0, 0, 0, 0, 0, 0, 0)
|
|
self.conn.mav.send(move_msg)
|
|
|
|
def move_pos_rel_yaw(self, x: float, y: float, z: float, yaw: float):
|
|
move_msg = MAVLink_set_position_target_local_ned_message(
|
|
int(perf_counter() * 1000),
|
|
self.conn.target_system,
|
|
self.conn.target_component,
|
|
MAV_FRAME_BODY_OFFSET_NED,
|
|
3576, x, y, z,
|
|
0, 0, 0, 0, 0, 0, yaw, 0)
|
|
self.conn.mav.send(move_msg)
|
|
|
|
def move_vel_rel_alt(self, vx: float, vy: float, z: float):
|
|
move_msg = MAVLink_set_position_target_local_ned_message(
|
|
int(perf_counter() * 1000),
|
|
self.conn.target_system,
|
|
self.conn.target_component,
|
|
MAV_FRAME_BODY_OFFSET_NED,
|
|
3555, 0, 0, z,
|
|
vx, vy, 0, 0, 0, 0, 0, 0)
|
|
self.conn.mav.send(move_msg)
|
|
|
|
def move_local_ned(self, x: float, y: float, z: float):
|
|
move_msg = MAVLink_set_position_target_local_ned_message(
|
|
int(perf_counter() * 1000),
|
|
self.conn.target_system,
|
|
self.conn.target_component,
|
|
MAV_FRAME_LOCAL_NED,
|
|
3576, x, y, z,
|
|
0, 0, 0, 0, 0, 0, 0, 0)
|
|
self.conn.mav.send(move_msg)
|
|
|
|
def distance(self, p1: Tuple, p2: Tuple) -> float:
|
|
return float(np.linalg.norm(np.asarray(p1) - np.asarray(p2)))
|
|
|
|
def wait_altitude(self, target_alt: float, tolerance: float = 0.5,
|
|
timeout: float = 30.0):
|
|
t0 = time.time()
|
|
while time.time() - t0 < timeout:
|
|
self.update_state()
|
|
if abs(self.altitude - target_alt) < tolerance:
|
|
print(f"\n[UAV] Altitude reached: {self.altitude:.1f}m")
|
|
return True
|
|
print(f"\r[UAV] Climbing: {self.altitude:.1f}m / {target_alt:.1f}m ",
|
|
end='', flush=True)
|
|
sleep(0.2)
|
|
return False
|
|
|
|
def configure_speed_limits(self,
|
|
wpnav_speed=DEFAULT_WPNAV_SPEED,
|
|
wpnav_accel=DEFAULT_WPNAV_ACCEL,
|
|
wpnav_speed_up=DEFAULT_WPNAV_SPEED_UP,
|
|
wpnav_speed_dn=DEFAULT_WPNAV_SPEED_DN,
|
|
wpnav_accel_z=DEFAULT_WPNAV_ACCEL_Z,
|
|
loit_speed=DEFAULT_LOIT_SPEED):
|
|
params = {
|
|
'WPNAV_SPEED': wpnav_speed,
|
|
'WPNAV_ACCEL': wpnav_accel,
|
|
'WPNAV_SPEED_UP': wpnav_speed_up,
|
|
'WPNAV_SPEED_DN': wpnav_speed_dn,
|
|
'WPNAV_ACCEL_Z': wpnav_accel_z,
|
|
'LOIT_SPEED': loit_speed,
|
|
}
|
|
for name, value in params.items():
|
|
self.set_param(name, value)
|
|
print(f"[UAV] Speed limits set: horiz={wpnav_speed}cm/s "
|
|
f"accel={wpnav_accel}cm/s² up={wpnav_speed_up}cm/s "
|
|
f"dn={wpnav_speed_dn}cm/s")
|
|
|
|
def set_max_velocity(self, speed):
|
|
self.conn.mav.command_long_send(
|
|
self.conn.target_system,
|
|
self.conn.target_component,
|
|
mavlink.MAV_CMD_DO_CHANGE_SPEED,
|
|
0, 1, speed, -1, 0, 0, 0, 0)
|