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

122 lines
4.0 KiB
Python

#!/usr/bin/env python3
"""SE(3) geometric control demo: circular trajectory tracking."""
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 import TrackingController, circle_trajectory
from uavcatch.control.factory import build_cascaded_controller
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="Geometric SE(3) circular trajectory tracking demo",
)
parser.add_argument("--environment", type=Path, default=ROOT / "environment.toml")
parser.add_argument("--uav", type=Path, default=ROOT / "uav.toml")
parser.add_argument("--control", type=Path, default=ROOT / "control.toml")
parser.add_argument("--radius", type=float, default=3.0, help="Circle radius (m)")
parser.add_argument("--height", type=float, default=2.0, help="Flight height (m)")
parser.add_argument("--period", type=float, default=12.0, help="One lap period (s)")
parser.add_argument("--duration", type=float, default=24.0, help="Simulation duration (s)")
parser.add_argument("--no-viewer", action="store_true")
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)
env_cfg.simulation.duration = args.duration
if control_cfg.controller.mode != "geometric":
print(
f"Warning: control.toml mode is '{control_cfg.controller.mode}'; "
"circle tracking works best with mode = \"geometric\"."
)
engine = SimEngine(env_cfg=env_cfg, uav_cfg=uav_cfg)
engine.reset()
center = np.array([0.0, 0.0, 0.0], dtype=np.float64)
pos_fn, yaw_fn, accel_fn = circle_trajectory(
center=center,
radius=args.radius,
height=args.height,
period=args.period,
)
init_target = pos_fn(0.0)
controller = build_cascaded_controller(
control_cfg,
engine.dynamics,
uav_cfg,
target_position=init_target.position,
target_velocity=init_target.velocity,
target_yaw=float(yaw_fn(0.0)),
)
tracker = TrackingController(
controller=controller,
target_fn=pos_fn,
yaw_fn=yaw_fn,
accel_fn=accel_fn,
)
duration = args.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:
def on_step(state: QuadState, t: float, omega: np.ndarray) -> None:
viewer.render_realtime(state, omega, dt, speed=args.speed)
if not viewer.is_running():
raise KeyboardInterrupt
try:
engine.run(tracker, duration=duration, on_step=on_step)
except KeyboardInterrupt:
pass
else:
engine.run(tracker, duration=duration)
final = engine.state
final_target = pos_fn(duration)
pos_err = float(np.linalg.norm(final.position - final_target.position))
print(f"Simulation finished at t={engine.time:.2f}s")
print(f"Final position: {final.position}")
print(f"Target position: {final_target.position}")
print(f"Position error: {pos_err:.4f} m")
print(f"Final velocity: {final.velocity}")
if pos_err > 1.0:
print("Warning: circular tracking error exceeds 1.0 m")
return 1
return 0
if __name__ == "__main__":
try:
raise SystemExit(main())
except KeyboardInterrupt:
raise SystemExit(130) from None