Files
uav-catch/examples/fly_hover.py
robinson e901fa9579 init
2026-07-03 04:19:44 +08:00

122 lines
3.2 KiB
Python

#!/usr/bin/env python3
"""Hover demo: PID-controlled single quadcopter in MuJoCo viewer."""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
import numpy as np
ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT / "src"))
from uavcatch.config import load_control_config, load_environment_config, load_uav_config
from uavcatch.control.pid import HoverPIDController
from uavcatch.dynamics.state import QuadState
from uavcatch.sim.engine import SimEngine
from uavcatch.sim.mujoco_viz import MuJoCoViewer
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Single-UAV hover simulation demo")
parser.add_argument(
"--environment",
type=Path,
default=ROOT / "environment.toml",
help="Path to environment.toml",
)
parser.add_argument(
"--uav",
type=Path,
default=ROOT / "uav.toml",
help="Path to uav.toml",
)
parser.add_argument(
"--control",
type=Path,
default=ROOT / "control.toml",
help="Path to control.toml",
)
parser.add_argument(
"--target-z",
type=float,
default=2.0,
help="Target hover height (m)",
)
parser.add_argument(
"--no-viewer",
action="store_true",
help="Run without MuJoCo viewer",
)
parser.add_argument(
"--speed",
type=float,
default=1.0,
help="Playback speed (1.0 = real time, <1 slow motion)",
)
return parser.parse_args()
def main() -> int:
args = parse_args()
env_cfg = load_environment_config(args.environment)
uav_cfg = load_uav_config(args.uav)
control_cfg = load_control_config(args.control)
engine = SimEngine(env_cfg=env_cfg, uav_cfg=uav_cfg)
engine.reset()
target = np.array([0.0, 0.0, args.target_z], dtype=np.float64)
controller = HoverPIDController(
dynamics=engine.dynamics,
uav_cfg=uav_cfg,
target_position=target,
control_cfg=control_cfg,
)
duration = env_cfg.simulation.duration
dt = env_cfg.simulation.dt
use_viewer = env_cfg.viewer.enable and not args.no_viewer
if use_viewer:
with MuJoCoViewer(env_cfg, uav_cfg=uav_cfg) as viewer:
last_state: QuadState | None = None
last_omega: np.ndarray | None = None
def on_step(state: QuadState, t: float, omega: np.ndarray) -> None:
nonlocal last_state, last_omega
last_state = state
last_omega = omega
viewer.render_realtime(state, omega, dt, speed=args.speed)
if not viewer.is_running():
raise KeyboardInterrupt
try:
history = engine.run(controller, duration=duration, on_step=on_step)
except KeyboardInterrupt:
history = []
else:
history = engine.run(controller, duration=duration)
final = engine.state
pos_err = np.linalg.norm(final.position - target)
print(f"Simulation finished at t={engine.time:.2f}s")
print(f"Final position: {final.position}")
print(f"Target position: {target}")
print(f"Position error: {pos_err:.4f} m")
print(f"Final velocity: {final.velocity}")
if pos_err > 0.1:
print("Warning: position error exceeds 0.1 m threshold")
return 1
return 0
if __name__ == "__main__":
try:
raise SystemExit(main())
except KeyboardInterrupt:
raise SystemExit(130) from None