From 8380b00955ae96b99033d4de2edcb37f93932cb5 Mon Sep 17 00:00:00 2001 From: Ezenbaatar Batjargal Date: Mon, 31 Aug 2026 11:38:36 -0700 Subject: [PATCH 1/4] feat(sim): MuJoCo + PX4 HIL mixed-embodiment simulator One MuJoCo world drives N PX4 SITL instances over the HIL interface (TCP 4560+i, sim owns PX4's clock via lockstep) alongside M simulated quadrupeds. Includes rotor aerodynamics (spool lag, H-force drag, ground effect, optional OU-gust wind), real Holybro X500 assets fetched on demand (gitignored), a one-command bring-up script with honest readiness gating, and sim-time-scaled PX4 parameters (battery, geofence, datalink-loss failsafe). Co-Authored-By: Claude Fable 5 --- .gitignore | 4 + dimos/simulation/px4_hil/README.md | 532 +++++++++ .../simulation/px4_hil/assets/x500/SOURCE.txt | 2 + dimos/simulation/px4_hil/fleet_bridge.py | 1023 +++++++++++++++++ dimos/simulation/px4_hil/hil_bridge.py | 432 +++++++ dimos/simulation/px4_hil/scene.py | 583 ++++++++++ dimos/simulation/px4_hil/sim.sh | 211 ++++ dimos/simulation/px4_hil/sim_params.py | 195 ++++ tools/fetch_x500_meshes.py | 117 ++ 9 files changed, 3099 insertions(+) create mode 100644 dimos/simulation/px4_hil/README.md create mode 100644 dimos/simulation/px4_hil/assets/x500/SOURCE.txt create mode 100644 dimos/simulation/px4_hil/fleet_bridge.py create mode 100644 dimos/simulation/px4_hil/hil_bridge.py create mode 100644 dimos/simulation/px4_hil/scene.py create mode 100755 dimos/simulation/px4_hil/sim.sh create mode 100644 dimos/simulation/px4_hil/sim_params.py create mode 100755 tools/fetch_x500_meshes.py diff --git a/.gitignore b/.gitignore index 1bc24524a4..401f9c07ae 100644 --- a/.gitignore +++ b/.gitignore @@ -97,3 +97,7 @@ yolo11n*.pt *.ignore *.ignore.* site/ + +# Fetched third-party meshes for the PX4 HIL simulator (tools/fetch_x500_meshes.py) +dimos/simulation/px4_hil/assets/x500/*.stl +dimos/simulation/px4_hil/assets/x500/*.obj diff --git a/dimos/simulation/px4_hil/README.md b/dimos/simulation/px4_hil/README.md new file mode 100644 index 0000000000..75c1f83cc0 --- /dev/null +++ b/dimos/simulation/px4_hil/README.md @@ -0,0 +1,532 @@ +# DimOS Mixed-Fleet Simulator + +Quadrotors running **real, unmodified PX4 firmware** and legged robots, sharing +**one MuJoCo physics world**, commanded through DimOS. + +``` +./dimos/simulation/px4_hil/sim.sh start 3 1 # 3 drones + 1 dog +dimos mcp call takeoff_all --arg altitude=6 +dimos mcp call rtl_all +./dimos/simulation/px4_hil/sim.sh stop +``` + +--- + +## Quickstart on a fresh machine + +Tested on Ubuntu x86_64 with Python 3.12 and PX4 v1.16.2. Everything below is +copy-pasteable; total setup is ~15 minutes plus downloads. + +**1. Clone and install** (skip the repo's large LFS assets -- this simulator +does not need them): + +```bash +GIT_LFS_SKIP_SMUDGE=1 git clone https://github.com/dimensionalOS/dimos.git +cd dimos +uv sync --extra sim --extra drone --extra cpu +echo "DIMOS_TRANSPORT=zenoh" > .env # the CLI tools read this +``` + +**2. Build PX4 SITL** (the real autopilot firmware the drones run): + +```bash +git clone --branch v1.16.2 --recursive https://github.com/PX4/PX4-Autopilot.git ~/PX4-Autopilot +make -C ~/PX4-Autopilot px4_sitl none_iris +# the build finishes by LAUNCHING one PX4 instance -- Ctrl-C it; sim.sh +# manages its own instances. A different checkout location works via +# PX4_DIR=/path sim.sh start ... +``` + +**3. Optional but recommended -- the good-looking assets.** Without them the +simulator falls back to primitive shapes that fly and walk identically: + +```bash +# real Holybro X500 drone meshes (BSD-3, fetched from PX4-gazebo-models): +.venv/bin/python tools/fetch_x500_meshes.py +# real Unitree Go1 + its trained walking policy (~2 GB Menagerie + git-lfs): +git lfs pull --include "data/.lfs/mujoco_sim.tar.gz" +.venv/bin/python -c "from mujoco_playground._src import mjx_env; mjx_env.ensure_menagerie_exists()" +``` + +**4. Fly:** + +```bash +./dimos/simulation/px4_hil/sim.sh start 2 1 --viewer # 2 drones + 1 dog +.venv/bin/dimos mcp call takeoff_all --arg altitude=6 +.venv/bin/dimos mcp call walk --arg speed_mps=0.5 +.venv/bin/dimos mcp call rtl_all +./dimos/simulation/px4_hil/sim.sh stop +``` + +The bridge logs which asset tier it chose at startup, and `sim.sh` refuses to +say "ready" unless every robot is actually connected and physics is stepping. +On laptops, set the CPU governor to performance first (see Troubleshooting) -- +it is worth ~2.4x realtime. + +--- + +## Why this exists + +DimOS previously ran two incompatible simulators: drones in **Gazebo** with PX4, +legged robots in **MuJoCo** inside the DimOS process. Nothing could interact +across that boundary, and on a normal laptop Gazebo could not fly three drones +at all — its physics loop is single-threaded, and starved PX4's sensor stream +until the EKF diverged and refused to arm. + +This replaces the drone half with MuJoCo while **keeping PX4**. That matters: +PX4 SITL is the same firmware that runs on a Pixhawk, so what flies here is what +flies on real hardware. Losing it would make the simulator a toy. + +| | Gazebo | This | +|---|---|---| +| 3 drones actually fly | ✗ | ✓ | +| 3 drones + 1 dog | ~1× realtime | **~10×** | +| drones and legged robots in one world | ✗ | ✓ | +| real PX4 firmware | ✓ | ✓ | +| real Unitree Go1 + trained policy | ✗ | ✓ | + +### Speed + +Physics-only ceilings, measured 2026-08-21 on the `performance` governor with +real X500 meshes and real Go1s. End-to-end throughput is far lower because PX4 +runs in lockstep as a separate process per drone — these are the headroom, not +the number you will see. + +| Fleet | Physics ceiling | µs/step | +|---|---|---| +| 1 drone | 313× | 12.8 | +| 2 drones | 192× | 20.9 | +| 3 drones | 162× | 24.8 | +| 5 drones | 109× | 36.6 | +| 1 dog | 82× | 48.5 | +| 2 dogs | 38× | 106.7 | +| 3 dogs | 22× | 181.0 | +| 3 drones + 1 dog | 61× | 65.5 | +| 3 drones + 2 dogs | 33× | 119.9 | +| 5 drones + 2 dogs | 30× | 133.3 | + +Two things dominate, and they are different for each robot class. **Drones are +PX4-bound**: each is a separate OS process in lockstep, so the gap between the +ceiling and reality is round-trip latency, not physics. **Legged robots are +physics-bound**: no PX4 at all, but 12 actuated joints with mesh inertias each. + +**The CPU governor is worth roughly 2.4×** and does not survive a reboot — see +Troubleshooting. Every number here assumes `performance`. + +**`--viewer` costs roughly 20-25%** at full mesh detail. Measured on 2 drones + +1 Go1: **7.7-10.4x headless**, **6.4-7.4x with the window open**. Both bounce +around by a couple of x between samples, so treat these as ranges rather than +figures. Drop the viewer for batch runs; the physics is identical either way. + +For reference, the same fleet ran ~10-12x before the full-detail X500 frame was +added -- 179k triangles per airframe is most of the difference, and it is all +rendering, not physics. + +--- + +## How it works + +``` + DimOS (namespaced: one module per robot) + SwarmCoordinator ── drone_state / swarm_cmd ── drone1..N , dog1..M + │ │ + │ MAVLink udp 14540+i UDP 15000+i + ▼ ▼ ▼ + ┌────────────────────────────────────────────────────────────────┐ + │ PX4 SITL × N (unmodified firmware, `none_iris` airframe) │ + └────────────────────────────────────────────────────────────────┘ + HIL / TCP 4560+i (PX4 is the client) + ┌────────────────────────────────────────────────────────────────┐ + │ fleet_bridge.py — ONE MuJoCo world │ + │ sensors out to all → actuator replies in → one physics step │ + └────────────────────────────────────────────────────────────────┘ +``` + +### The PX4 link (HIL) + +Gazebo talks to PX4 through `gz_bridge`, a module compiled into the firmware. +MuJoCo has no equivalent, so we use PX4's other supported path: the generic +**HIL** (Hardware-In-The-Loop) interface `simulator_mavlink`, the same one +jMAVSim and JSBSim use. + +* **TCP on 4560 + i**, and *PX4 is the client* — the bridge must be listening first. +* Sim → PX4: `HIL_SENSOR` every step (accel, gyro, mag, baro), `HIL_GPS` at 10 Hz. +* PX4 → Sim: `HIL_ACTUATOR_CONTROLS`, normalised motor outputs. + +**The simulator owns PX4's clock.** Every `HIL_SENSOR` calls +`px4_clock_settime()` inside PX4, so PX4 never runs on wall-clock time — it +advances exactly as fast as we feed it. A slow machine makes the world run +*slower*, it does not starve the sensor stream. That single property is why this +scales past the one flyable drone Gazebo manages here. + +### Namespaces + +Every robot is one module instance under its own namespace, so it gets its own +RPC surface, topics, TF frames and config keys: + +``` +drone2/px4dronemodule/takeoff RPC and MCP tool +/drone2/odom topic +-o drone2/px4dronemodule.max_altitude_m=15 +``` + +Only two streams are **exposed** (left global) so data crosses the boundary: +`drone_state` (every robot publishes its own snapshot) and `swarm_cmd` (the +coordinator broadcasts; each robot acts only on its own key or `all`). +Everything else — notably `cmd_vel` — stays namespace-local, so one drone's +vision tracker can never fly another drone. + +Drones and legged robots publish the **same** state message with a +`robot_class` field, which is what lets `fleet_state` and `count_within` cover +both without special-casing either. + +### Frames + +The classic way to break a bridge like this. + +| | Body | World | +|---|---|---| +| MuJoCo | **FLU** (Forward-Left-Up) | **NWU** (North-West-Up) | +| PX4 | **FRD** (Forward-Right-Down) | **NED** (North-East-Down) | + +Both conversions are `(x, -y, -z)`. NWU rather than the more common ENU is +deliberate: PX4 assumes a body at zero yaw points **North**, and a body at +identity in MuJoCo points along world +x. An ENU world puts "forward" at East, +which yaws the magnetometer 90° and diverges the attitude estimate. + +--- + +## Running it + +`sim.sh` handles the ordering that otherwise has to be remembered, and waits +until every robot has actually reported before saying `ready`. + +```bash +SIM=./dimos/simulation/px4_hil/sim.sh + +$SIM start 0 2 # 2 dogs, no drones +$SIM start 0 3 # 3 dogs +$SIM start 2 0 # 2 drones, no dogs +$SIM start 1 1 # 1 drone + 1 dog +$SIM start 3 1 # the full demo fleet + +$SIM status +$SIM log bridge # or: daemon, px4_0 +$SIM stop +``` + +### Seeing it + +The simulator runs **headless by default** — no window opens. Add `--viewer` +(anywhere in the arguments) to get the MuJoCo window: + +```bash +$SIM start 3 1 --viewer +``` + +In the window: **Tab** cycles cameras (free -> `field` -> a chase cam on each +vehicle), **double-click** a vehicle then drag to orbit it, **S** toggles +shadows, **space** pauses. Closing the window now just drops the sim to +headless — physics, PX4 and the fleet keep running; use `$SIM stop` to shut +down. + +Each drone carries an accent colour (drone1 red, drone2 cyan, drone3 amber) on +its beacon and stripe, and lands on a matching ringed pad, so you can tell which +vehicle is flying which lane of a sweep. Those markings are visual only +(`contype=0`, zero mass) and cannot affect the physics — verified by diffing the +compiled model against the unstyled scene: identical masses, inertias, DOF, and +collidable-geom counts, and under 5% on step cost. + +The ground checker is 2 m per square; at a typical 8 m sweep altitude that is +what gives you usable motion parallax and a way to eyeball distance. + +### Commands + +Fleet-wide (no namespace — they act on the whole fleet): + +```bash +dimos mcp call list_drones +dimos mcp call preflight_check +dimos mcp call fleet_state +dimos mcp call takeoff_all --arg altitude=6 +dimos mcp call grid_sweep --arg corner_b_north=40 --arg corner_b_east=30 +dimos mcp call investigate --arg north=30 --arg east=10 --arg num_drones=2 +dimos mcp call line_formation --arg altitude=7 --arg spacing_m=6 +dimos mcp call count_within --arg drone=drone1 --arg radius_m=100 +dimos mcp call goto_drone --arg drone=drone1 --arg north=25 --arg east=0 --arg altitude=6 +dimos mcp call rtl_all +dimos mcp call emergency_land_all # controlled descent +dimos mcp call kill_all # last resort: motors off, they fall +``` + +Per robot. A skill offered by only one robot keeps its bare name; when several +robots offer it the name is qualified. Note `state`, `stop` and `goto` exist on +*both* robot classes, so in a mixed fleet they are always qualified: + +```bash +dimos mcp call crouch # only one dog -> bare +dimos mcp call dog2/leggedsimmodule/crouch # two or more dogs -> qualified +dimos mcp call drone2/px4dronemodule/takeoff --arg altitude=5 +dimos mcp call dog1/leggedsimmodule/state # `state` exists on both classes +dimos mcp list-tools | grep '"name"' # see what exists right now +``` + +Legged robots: + +```bash +dimos mcp call walk --arg speed_mps=0.3 --arg turn_rate_rads=0.1 # + turn = right +dimos mcp call halt # stop walking (NOT `stop` -- see below) +dimos mcp call crouch # and stand +``` + +Legged robots are **real Unitree Go1s driven by DimOS's trained ONNX policy** +when the assets are present (see below). Measured envelope of that policy in +this world: **0.94 m/s** forward (1.0 commanded), **0.31 m/s** reverse, +**0.55 rad/s** yaw, and it holds a turn and a cruise at the same time. + +**It has a low-speed deadband: a commanded 0.18 m/s produces no motion at all.** +0.30 walks, 0.45 walks well. Anything under ~0.25 m/s is a stop, not a slow +walk — this is why `goto`'s speeds are set where they are, and it is the first +thing to check if a legged robot ignores a command. + +`crouch` / `set_height` do nothing on a real Go1: the shipped policy is a +flat-ground velocity controller with no height input. The bridge logs that +rather than silently accepting the command. They still work on the primitive +fallback below. + +### Legged assets + +Without MuJoCo Menagerie and the trained policy, dogs fall back to a +primitive box-and-capsule quadruped with a hand-written trot (0.37 m/s, +0.12 rad/s, and it falls at closed-loop reversals). To get the real robot: + +```bash +sudo apt install git-lfs && git lfs install +git lfs pull --include "data/.lfs/mujoco_sim.tar.gz" # the trained policy +.venv/bin/python -c "from mujoco_playground._src import mjx_env; mjx_env.ensure_menagerie_exists()" +``` + +Menagerie is a plain 2 GB `git clone` (no LFS, no sudo) and supplies the mesh; +the LFS archive supplies `unitree_go1_policy.onnx`. **Both are required** — +the mesh alone gives a robot that cannot walk. The bridge logs which backend it +chose on startup (`legged robots: real Unitree Go1 + trained policy xN`). + +### The operating envelope + +**5 drones + 5 dogs, maximum.** `sim.sh` refuses more. That is a policy, not a +performance limit -- the validated demo ladder is 2 drones, then 2 dogs, then +2 drones + 1 dog, and everything beyond 5+5 is unvalidated territory. + +### The staged mission + +The one-command version of "the drones go in first, then the dog": + +```bash +dimos mcp call takeoff_all --arg altitude=6 +dimos mcp call sweep_then_ground --arg ground_delay_s=30 +dimos mcp call mission_status # [air sweep] / [ground transect] / [complete] +dimos mcp call rtl_all # when done -- drones hold at lane ends +``` + +Drones sweep the marked 40 x 30 m field immediately; after `ground_delay_s` +(wall-clock) the dog walks the centre transect -- entry, centre, far edge -- +and returns to where it started. `emergency_land_all` and `kill_all` abort the +ground stage. The delay is wall-clock on purpose: the operator watches the wall +clock, not sim time. + +### The world + +The 40 x 30 m sweep field is marked with white lines (NED (0,0)..(40,30)) so a +grid sweep is legible in the viewer. Trees are visual-only and parked well away +from every mission path -- a trunk the dog could walk through must never sit +where a route goes. The one collidable obstacle cluster (three crates, NED +(12,-14), west of the field) exists so the world has at least one honest +obstacle; no default route goes through it. + +### Natural-language control + +```bash +export OPENAI_API_KEY=sk-... # required BEFORE start; the daemon dies without it +SIM_BLUEPRINT=mixed-fleet-agentic ./dimos/simulation/px4_hil/sim.sh start 0 2 --viewer +dimos humancli +``` + +**Run `humancli` from the repo root** (`~/Work/dimos`) -- the `.env` there +pins `DIMOS_TRANSPORT=zenoh` for every process. Without it, humancli on Linux +defaults to LCM, connects to nobody, and shows "thinking..." forever while the +daemon (on zenoh) never hears it. + +**Reopen `humancli` after every `sim.sh start`.** A humancli window from a +previous run still says "Connected" but talks to a dead session -- its +messages go nowhere and it shows "thinking..." forever. Ctrl-C, run +`dimos humancli` again, done. (`dimos agent-send "..."` is the quick way to +test the agent without a chat window.) + +### Wind + +```bash +SIM_WIND_N=3 SIM_WIND_E=0 SIM_GUST_STD=1 $SIM start 2 0 +``` + +Constant mean wind (m/s, world NED) plus an Ornstein-Uhlenbeck gust process, +entering through the airspeed the rotor drag sees. Zero by default. Verified: +at 3 m/s mean + 1 m/s gusts the fleet holds position within ~0.1 m and lands +clean. The drone side also models motor spool lag, rotor H-force drag and +ground effect -- constants and their approximations are documented at the top +of `hil_bridge.py`. + +### Control boundaries + +Run `dimos mcp call boundaries` to see every active limit. Two layers, catching +different failures: + +| Boundary | Layer | What it catches | +|---|---|---| +| Operating radius 250 m | dispatch-time (coordinator + dog `goto`) | a bad command, before anything moves | +| Altitude ceiling | dispatch-time | same | +| Separation floor 2 m | dispatch-time admission | a waypoint ending too close to another aircraft | +| PX4 geofence (hold at 260 m / 45 m) | **in-flight, inside the autopilot** | `set_velocity` drift and anything dispatch cannot see (aircraft only) | +| Ground fence (halt at 250 m) | **in-motion, in each legged module** | a `walk` command carrying a dog out of the operating area -- the ground analogue of the PX4 fence. Walking back in is allowed; continued escape re-trips it | +| Ground proximity (halt at 0.6 m) | **in-motion, in the simulator** | two ground robots converging -- including ones already moving, which no dispatch check can see. Re-arms at 1.2 m | + +Legged robots have two motion commands with deliberately different semantics: +`move(forward_m, right_m)` is body-relative and **holds the heading** (the Go1 +strafes), so "left 5 then forward 2" composes the way a person means it; +`goto(north, east)` turns toward its target -- efficient for distance, but it +changes what "forward" means afterwards. Both are closed-loop **inside the +simulator at physics rate**: control from the DimOS side on wall-clock ticks +acts on 6-12 sim-seconds of stale state at these realtime factors and is +unstable (measured: a 5 m strafe wandered 140 m). +| Airborne gate | dispatch-time | OFFBOARD to a grounded drone (trips a failsafe that blocks the next arm) | +| Datalink-loss failsafe (`NAV_DLL_ACT=2`) | **in-flight, inside the autopilot** | the laptop/DimOS dying mid-flight -- each drone returns and lands on its own authority | + +With the datalink-loss failsafe configured, PX4 refuses to arm without a live +ground station -- so the SwarmCoordinator **is** the fleet's ground station: it +binds the shared GCS port (14550, the QGroundControl convention every PX4 +instance targets), discovers each vehicle from its stream, and answers with +2 Hz heartbeats. Kill DimOS mid-flight and those heartbeats stop; every drone +holds 5 s, returns, lands and disarms entirely on its own authority — the +exact behaviour wanted on real hardware, exercised in sim. (Running a real +QGroundControl instead? Set `-o swarmcoordinator.gcs_port=0` so they do not +fight over the port.) + +The fence radius (260 m) is deliberately just outside the dispatch radius +(250 m) so the polite refusal always fires first; the fence is the backstop, +not the interface. The separation floor is admission control, **not** in-flight +collision avoidance — two drones already converging are not re-checked. + +**Order matters: `takeoff_all` before any position maneuver.** `grid_sweep`, +`line_formation`, `investigate` and `goto_drone` engage OFFBOARD; doing that to +a grounded, disarmed vehicle trips a PX4 failsafe that then blocks the *next* +arm. The simulator refuses and tells you, rather than failing later. + +--- + +## Suggested first session + +```bash +SIM=./dimos/simulation/px4_hil/sim.sh +D=.venv/bin/dimos + +$SIM start 0 2 # two dogs +$D mcp call list_drones +$D mcp call dog1/leggedsimmodule/crouch +$D mcp call dog1/leggedsimmodule/stand + +$SIM start 0 3 # three dogs +$D mcp call fleet_state + +$SIM start 2 0 # two drones +$D mcp call preflight_check +$D mcp call takeoff_all --arg altitude=5 +$D mcp call fleet_state +$D mcp call rtl_all + +$SIM start 1 1 # one of each +$D mcp call takeoff_all --arg altitude=5 +$D mcp call count_within --arg drone=drone1 --arg radius_m=100 # sees the dog +$D mcp call crouch +$D mcp call rtl_all +$SIM stop +``` + +--- + +## Troubleshooting + +**Drones will not arm; PX4 logs `High Accelerometer Bias` or `ekf2 missing data`.** +Check the CPU governor first — it does **not** survive a reboot: + +```bash +cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor # want: performance +sudo cpupower frequency-set -g performance +``` + +On `powersave` this machine pins every core near 1.5 GHz against a 4.7 GHz +capability, even on AC with cool temperatures. MuJoCo tolerates it far better +than Gazebo did — 3 drones plus a dog still run at ~3.3× — but performance mode +is free speed. + +**`REJECTED: no aircraft are armed`.** Working as intended. Call `takeoff_all` +first; see the ordering note above. + +**Nothing takes off after an RTL.** Fixed — `takeoff` now clears the latched +`AUTO.RTL` nav state before arming. If you see it again, `$SIM stop && $SIM start`. + +**Drone flips immediately on arming.** A rotor yaw-torque sign is inverted. +PX4 computes `moment = ct * position.cross(axis) - ct * km * axis` with +`axis = (0,0,-1)` in FRD, giving `torque_z(FRD) = +km * thrust`; FRD z points +down, so that is *negative* about MuJoCo's z. See the comment in `scene.py`. + +**Weird estimator behaviour after fiddling with parameters.** PX4 transports +INT32 parameters bit-cast into a float field — sending a literal `1.0` stores +`1065353216`. `sim.sh` wipes each instance's `parameters.bson` on start for +exactly this reason. + +**Battery dies about a minute in.** `SIM_BAT_DRAIN` defaults to 60 s full to +empty. `sim.sh` runs `sim_params.py` automatically; check +`/tmp/dimos-sim/params.log` if a run ends early. + +--- + +## Files + +| File | What | +|---|---| +| `sim.sh` | Bring the whole stack up and down | +| `fleet_bridge.py` | The simulator: one MuJoCo world, N PX4 links, M legged links | +| `hil_bridge.py` | HIL protocol, frame conversions, and one vehicle's link | +| `scene.py` | Builds the world — single source of rotor geometry and the quadruped | +| `../../robot/drone/px4_drone_module.py` | DimOS module for one drone | +| `../../robot/legged/legged_sim_module.py` | DimOS module for one legged robot | +| `../../robot/drone/px4_swarm_coordinator.py` | Fleet state, guardrails, maneuvers | + +--- + +## Known limits + +* **The quadruped is solid with the trained policy, weak without it.** With a + real Go1, `goto` reaches **6 of 6** spread targets (worst error 0.62 m) and + the robot did not fall in any tested command, including reversals and targets + directly behind it. On the primitive fallback it is the old hand-tuned trot: + 4 of 6 targets, and it topples on large sustained turns. +* **The Go1 policy is Go1-only.** DimOS ships trained policies for Go1 and G1. + Menagerie also contains a Go2 mesh, but there is no Go2 policy — so this is a + good simulated demo, not sim-to-real transfer for Go2 hardware. +* **A real Go1 costs about 2x the physics time** of the primitive dog per step + (measured): fewer contacts, but 12 actuated joints with mesh inertias. At + 3 drones + 2 dogs the physics ceiling is ~53x, still far above what PX4 + lockstep delivers, so it is affordable. +* **The legged "stop walking" tool is `halt`, not `stop`.** `Module.stop()` is + the framework's teardown RPC — it closes the module's RPC, tools and event + loop. A `@skill` named `stop` shadowed it, so the module never tore down, and + `_on_swarm_cmd` calling `self.stop()` on a `land`/`hold` broadcast would have + killed the module outright once the shadowing was removed. Never add a skill + named `stop` to a Module subclass. +* **No aerodynamics.** MuJoCo models no rotor drag, downwash or ground effect. + Fine for separation and coordination work; not a flight-dynamics model. +* **The separation guardrail is a dispatch-time admission check**, not + collision avoidance. It rejects a commanded waypoint that would end too close; + it does not re-check two drones already converging in flight. +* **`SIM_DRONES` / `SIM_DOGS` must match** what the bridge was started with, or + DimOS waits for robots that do not exist. `sim.sh` keeps them in step. diff --git a/dimos/simulation/px4_hil/assets/x500/SOURCE.txt b/dimos/simulation/px4_hil/assets/x500/SOURCE.txt new file mode 100644 index 0000000000..8f20302976 --- /dev/null +++ b/dimos/simulation/px4_hil/assets/x500/SOURCE.txt @@ -0,0 +1,2 @@ +X500 meshes: PX4/PX4-gazebo-models, BSD-3-Clause. +Fetched and decimated by tools/fetch_x500_meshes.py -- do not edit by hand. diff --git a/dimos/simulation/px4_hil/fleet_bridge.py b/dimos/simulation/px4_hil/fleet_bridge.py new file mode 100644 index 0000000000..b50d41a68d --- /dev/null +++ b/dimos/simulation/px4_hil/fleet_bridge.py @@ -0,0 +1,1023 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""One MuJoCo world driving N PX4 instances and M legged robots in lockstep. + +This is the runnable simulator. The drones and the legged robots share a single +physics world -- the entire reason for moving off Gazebo, where drones lived in +gz and legged robots in a separate MuJoCo process with only positions passed +between them. + +Lockstep with N autopilots +-------------------------- +Each PX4 advances its own clock from the ``HIL_SENSOR`` stream we send it, so +all N must be driven together or they desynchronise: + + 1. send HIL_SENSOR (and periodically HIL_GPS) to every vehicle + 2. wait for every vehicle's HIL_ACTUATOR_CONTROLS reply, all at once + 3. apply the actuator commands + 4. step the shared physics exactly once + +Step 2 waits on every socket concurrently with ``select``. Reading each vehicle +in turn with a blocking recv costs one scheduler round-trip *per vehicle per +step*, which dominated the loop and held the fleet to a fraction of the speed +the physics can sustain. + +Ports, per drone i (loopback only, no external networking): + 4560 + i HIL/TCP, we listen and PX4 dials in + 14540 + i PX4's MAVLink offboard channel, where DimOS connects + 14550 + i PX4's MAVLink GCS channel, for QGroundControl + +Legged robots get a UDP endpoint each at ``LEGGED_PORT_BASE + i``, deliberately +mirroring the drone pattern: the bridge is the simulator, and DimOS modules are +clients that reach it over a socket. That keeps the DimOS side identical in +shape for both robot classes, so both namespace the same way. + +Usage:: + + python -m dimos.simulation.px4_hil.fleet_bridge --drones 3 --dogs 1 + # then once per drone: + cd ~/PX4-Autopilot/build/px4_sitl_default && PX4_SIM_MODEL=none_iris ./bin/px4 -i 0 -d +""" + +from __future__ import annotations + +import argparse +import json +import math +import select +import socket +import time +from typing import Any + +import numpy as np + +from dimos.simulation.px4_hil.hil_bridge import ( + ACTUATOR_WAIT_S, + DEFAULT_ORIGIN_ALT, + DEFAULT_ORIGIN_LAT, + DEFAULT_ORIGIN_LON, + GPS_RATE_HZ, + DroneLink, +) +from dimos.simulation.px4_hil.scene import ( + DOG_TORSO_Z, + GO1_HOME_ANGLES, + GO1_SPAWN_Z, + build_model, + go1_policy_path, + go1_unavailable_reason, + x500_meshes_available, +) +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +# UDP base port for legged robots. Chosen clear of PX4's 4560/145xx/185xx ranges. +LEGGED_PORT_BASE = 15000 +# Legged state publish rate. Matches Px4DroneModule's telemetry cadence so the +# fleet-wide picture updates uniformly across robot classes. +LEGGED_STATE_HZ = 10.0 + + +# --- Legged gait ------------------------------------------------------------ +# Link lengths from scene.py. Both segments are 0.20 m, so the leg is +# symmetric and the IK below has a clean closed form. +THIGH_LEN_M = 0.20 +CALF_LEN_M = 0.20 +# Nominal foot height below the hip at the standing pose, i.e. fk(0.3, -0.6). +STANCE_Z_M = -0.382 +# Torso height below which the robot is considered down. An open-loop trot on +# 2-DOF legs can topple; without detection it then drags along the ground +# flailing, and every later command silently does nothing. Detecting it lets the +# gait stop and the robot settle, and lets DimOS report `fallen` honestly rather +# than claiming to be walking. +FALLEN_HEIGHT_M = 0.25 +# Stride at full commanded speed, and how high a swinging foot clears ground. +MAX_STRIDE_M = 0.12 +SWING_LIFT_M = 0.05 +GAIT_FREQ_HZ = 2.0 +# Measured travel per gait cycle, in strides. A cycle advances further than one +# stride because both diagonal pairs contribute; calibrated by walking the model +# and dividing distance by (stride * frequency * time). Without this a commanded +# speed and the achieved speed differ by ~50%, which makes closed-loop `goto` +# overshoot. +STRIDES_PER_CYCLE = 1.54 +MAX_BODY_SPEED_MS = MAX_STRIDE_M * GAIT_FREQ_HZ * STRIDES_PER_CYCLE # ~0.37 +# Walking backwards at full stride topples this open-loop gait: the swing foot +# catches and the body pitches over. Half stride is stable, so backwards is +# capped rather than left as a way to fall over. +MAX_BACKWARD_SPEED_MS = MAX_BODY_SPEED_MS * 0.5 +# Reversing covers about twice the ground per stride that walking forward does +# -- the gait is not symmetric, because the swing and stance arcs are not mirror +# images about the standing pose. Measured, not derived. +BACKWARD_SPEED_RATIO = 2.0 +# Turn authority is the honest limit of this airframe: 2 DOF per leg and no hip +# abduction, so yaw comes only from a differential stride. Above 0.45 the robot +# topples turning in place. Turning while walking is meaningfully better +# (~12 deg/s) than turning on the spot (~4 deg/s). +TURN_STRIDE_FRACTION = 0.45 +# Kept deliberately inside the stable envelope rather than at its edge. Sweeping +# speed against turn rate shows the failures are non-monotonic -- (0.30, 0.20) +# topples while (0.37, 0.20) survives -- which means the boundary is a resonance, +# not a clean limit. Margin is the right answer, not a tighter fit. +MAX_TURN_RATE_RADS = 0.12 +# How much top speed is given up at full turn rate. 0.4 leaves ~0.22 m/s, which +# is the fastest combination measured stable over a sustained turn. +TURN_SPEED_DERATE = 0.4 +# Commands ramp rather than snap. A step change from full forward to full +# reverse topples any legged robot, real or simulated; this is the controller +# being physically honest, not a workaround. +CMD_SLEW_MS2 = 0.6 +# Yaw slews far more gently than speed. Reversing the differential stride is +# what actually topples this gait: a closed-loop controller crossing zero +# heading error flips the command in a fraction of a second, which catches a +# foot mid-swing. At 0.3 rad/s^2 a full reversal takes about a second. +CMD_SLEW_RADS2 = 0.3 +# Legs, with their trot phase offset and which side they are on. Diagonal pairs +# move together, which is what keeps a quadruped statically balanced mid-gait. +LEGS = ( + ("fl", 0.0, +1), + ("rr", 0.0, -1), + ("fr", 0.5, -1), + ("rl", 0.5, +1), +) + + +def _leg_ik(x: float, z: float) -> tuple[float, float]: + """Foot position in the leg's sagittal plane -> (hip, knee) angles. + + x is forward, z is up (negative below the hip). The knee is taken as the + negative solution because scene.py ranges it to [-2.4, 0], which keeps the + joint bending the same way a real quadruped's does. + """ + r2 = x * x + z * z + cos_knee = (r2 - THIGH_LEN_M**2 - CALF_LEN_M**2) / (2 * THIGH_LEN_M * CALF_LEN_M) + knee = -math.acos(max(-1.0, min(1.0, cos_knee))) + hip = math.atan2(-x, -z) - math.atan2( + CALF_LEN_M * math.sin(knee), THIGH_LEN_M + CALF_LEN_M * math.cos(knee) + ) + return hip, knee + + +class LeggedLink: + """One legged robot: UDP control endpoint plus a trot gait controller. + + The gait is deliberately open-loop and simple -- a diagonal trot with a + differential stride for turning. It is enough to make the robot a real fleet + member that can be sent somewhere, which is the point; a learned or + model-predictive controller belongs in a DimOS module, not in the physics + bridge. + """ + + def __init__(self, mj: Any, model: Any, data: Any, index: int, host: str, port: int) -> None: + self.index = index + self.name = f"dog{index}" + self._mj = mj + self.model = model + self.data = data + self.host = host + self.port = port + + n2i = mj.mj_name2id + self.body = n2i(model, mj.mjtObj.mjOBJ_BODY, f"{self.name}_torso") + if self.body < 0: + raise ValueError(f"scene has no torso body for {self.name}") + jid = n2i(model, mj.mjtObj.mjOBJ_JOINT, f"{self.name}_root") + self.dof_adr = int(model.jnt_dofadr[jid]) + self.qpos_adr = int(model.jnt_qposadr[jid]) + + # Resolve actuators per leg, so gait phases map to the right joints. + self.legs: list[tuple[str, float, int, int, int]] = [] + for leg, phase, side in LEGS: + hip = n2i(model, mj.mjtObj.mjOBJ_ACTUATOR, f"{self.name}_{leg}_hip") + knee = n2i(model, mj.mjtObj.mjOBJ_ACTUATOR, f"{self.name}_{leg}_knee") + if hip < 0 or knee < 0: + raise ValueError(f"scene has no {leg} actuators for {self.name}") + self.legs.append((leg, phase, side, hip, knee)) + + self.stance_scale = 1.0 # 1.0 = standing, lower = crouched + # Commanded (target) and applied (slewed) body velocities. + self.cmd_vx = 0.0 # m/s forward + self.cmd_wz = 0.0 # rad/s yaw, positive = turn RIGHT (NED) + self._vx = 0.0 + self._wz = 0.0 + self._phase = 0.0 + self.fallen = False + + self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self.sock.bind((host, port)) + self.sock.setblocking(False) + self.peer: tuple[str, int] | None = None + + # -- commands ---------------------------------------------------------- + + def poll_commands(self) -> None: + """Apply any waiting command datagrams. Never blocks.""" + while True: + try: + payload, addr = self.sock.recvfrom(4096) + except (BlockingIOError, InterruptedError, OSError): + return + self.peer = addr + try: + cmd = json.loads(payload) + except ValueError: + logger.warning(f"[{self.name}] malformed command datagram") + continue + action = cmd.get("action", "") + if action == "recover": + # Clear the latch and drive the legs to full stance. If the + # robot is on its side the servos can often push it back over; + # if it is fully inverted they cannot, and the latch simply + # re-arms next step. Either way the operator is not stuck with a + # robot that ignores every command for the rest of the session. + self.fallen = False + self.stance_scale = 1.0 + self.cmd_vx = self.cmd_wz = 0.0 + logger.info(f"[{self.name}] recovery attempt") + elif action == "stance": + h = float(cmd.get("height", 1.0)) + self.stance_scale = max(0.35, min(1.0, h)) + elif action == "move_to": + # Target arrives in NED; the bridge world is NWU (y = -east, + # yaw sign flips). + tn = float(cmd.get("north", 0.0)) + te = float(cmd.get("east", 0.0)) + yaw_ned = float(cmd.get("yaw", 0.0)) + self._move_target = (tn, -te, -yaw_ned) + logger.info(f"[{self.name}] move_to NED ({tn:.1f}, {te:.1f}), heading held") + elif action == "walk": + self.cmd_vx = max( + -MAX_BACKWARD_SPEED_MS, min(MAX_BODY_SPEED_MS, float(cmd.get("vx", 0.0))) + ) + self.cmd_wz = max( + -MAX_TURN_RATE_RADS, min(MAX_TURN_RATE_RADS, float(cmd.get("wz", 0.0))) + ) + elif action == "stop": + self.cmd_vx = 0.0 + self.cmd_wz = 0.0 + elif action == "move_to": + logger.warning( + f"[{self.name}] move_to needs the Go1 policy (strafe); " + "the primitive gait cannot hold heading -- use goto" + ) + elif action == "noop": + pass + + # -- gait -------------------------------------------------------------- + + def apply(self, dt: float) -> None: + """Advance the gait and write joint targets.""" + # Fall detection. Once down, stop trying to walk: continuing to cycle + # the legs just thrashes and never recovers. Holding the stance pose at + # least lets the body settle, and `stand` re-arms the gait. + height = float(self.data.xpos[self.body][2]) + if height < FALLEN_HEIGHT_M: + if not self.fallen: + logger.warning(f"[{self.name}] fallen (h={height:.2f} m); gait stopped") + self.fallen = True + self.cmd_vx = self.cmd_wz = 0.0 + self._vx = self._wz = 0.0 + elif self.fallen and height > FALLEN_HEIGHT_M + 0.08: + logger.info(f"[{self.name}] upright again (h={height:.2f} m)") + self.fallen = False + + # Slew towards the command so reversals ramp instead of snapping. + self._vx += max(-CMD_SLEW_MS2 * dt, min(CMD_SLEW_MS2 * dt, self.cmd_vx - self._vx)) + self._wz += max(-CMD_SLEW_RADS2 * dt, min(CMD_SLEW_RADS2 * dt, self.cmd_wz - self._wz)) + + # Back off speed while turning. Top speed is stable in a straight line + # and stable turning at cruise, but sustained full-speed turning topples + # the robot -- the same trade any vehicle makes in a corner. + turn_frac = abs(self._wz) / MAX_TURN_RATE_RADS if MAX_TURN_RATE_RADS else 0.0 + speed_limit = MAX_BODY_SPEED_MS * (1.0 - TURN_SPEED_DERATE * min(1.0, turn_frac)) + self._vx = max(-speed_limit, min(speed_limit, self._vx)) + + stance_z = STANCE_Z_M * self.stance_scale + moving = abs(self._vx) > 1e-3 or abs(self._wz) > 1e-3 + + if not moving: + # Park in the standing pose rather than freezing mid-swing, which + # would leave a foot in the air and topple the robot. + self._phase = 0.0 + hip, knee = _leg_ik(0.0, stance_z) + for _leg, _ph, _side, hip_id, knee_id in self.legs: + self.data.ctrl[hip_id] = hip + self.data.ctrl[knee_id] = knee + return + + self._phase = (self._phase + GAIT_FREQ_HZ * dt) % 1.0 + stride = self._vx / (GAIT_FREQ_HZ * STRIDES_PER_CYCLE) + if stride < 0.0: + stride /= BACKWARD_SPEED_RATIO + stride = max(-MAX_STRIDE_M, min(MAX_STRIDE_M, stride)) + # Turning is a differential stride: the outside legs take longer steps. + turn = MAX_STRIDE_M * TURN_STRIDE_FRACTION * (self._wz / MAX_TURN_RATE_RADS) + + for _leg, phase_off, side, hip_id, knee_id in self.legs: + # Yaw follows the NED convention the rest of the fleet uses: + # positive is nose-RIGHT. Turning right means the left (outside) + # legs take the longer stride, so side (+1 for left) adds. + leg_stride = stride + side * turn + p = (self._phase + phase_off) % 1.0 + if p < 0.5: # stance: foot travels backward, driving the body forward + u = p / 0.5 + x = leg_stride * (0.5 - u) + z = stance_z + else: # swing: foot returns forward, lifted clear of the ground + u = (p - 0.5) / 0.5 + x = leg_stride * (u - 0.5) + z = stance_z + SWING_LIFT_M * math.sin(math.pi * u) + hip, knee = _leg_ik(x, z) + self.data.ctrl[hip_id] = hip + self.data.ctrl[knee_id] = knee + + # -- state ------------------------------------------------------------- + + def _yaw(self) -> float: + """Body yaw in the world NWU frame, radians.""" + qw, qx, qy, qz = self.data.qpos[self.qpos_adr + 3 : self.qpos_adr + 7] + return math.atan2(2.0 * (qw * qz + qx * qy), 1.0 - 2.0 * (qy * qy + qz * qz)) + + def state(self, sim_time_us: int) -> dict[str, Any]: + pos = self.data.xpos[self.body] + vel = self.data.qvel[self.dof_adr : self.dof_adr + 3] + return { + "key": self.name, + "robot_class": "legged", + "connected": True, + "nwu": [float(pos[0]), float(pos[1]), float(pos[2])], + "velocity": [float(vel[0]), float(vel[1]), float(vel[2])], + "yaw_rad": self._yaw(), + "height_m": float(pos[2]), + "nominal_height_m": DOG_TORSO_Z, + "walking": abs(self._vx) > 1e-3 or abs(self._wz) > 1e-3, + "fallen": self.fallen, + "speed_mps": self._vx, + "turn_rate_rads": self._wz, + "sim_time_s": sim_time_us / 1e6, + } + + def publish(self, sim_time_us: int) -> None: + if self.peer is None: + return + try: + self.sock.sendto(json.dumps(self.state(sim_time_us)).encode(), self.peer) + except OSError: + pass + + def close(self) -> None: + try: + self.sock.close() + except OSError: + pass + + +# --------------------------------------------------------------------------- +# Real Unitree Go1 driven by DimOS's trained locomotion policy +# --------------------------------------------------------------------------- +# Measured envelope of the shipped policy under this world's solver +# (implicitfast, 100 iterations -- NOT the Euler/1-iteration setup it was +# trained under; it transfers fine, verified by probe): +# forward 1.0 commanded -> 0.94 m/s +# reverse -0.4 -> -0.31 m/s +# yaw +-0.8 -> +-0.55 rad/s, and it holds an arc while moving +# never fell in any tested command, including reversals +# Commands beyond that DEGRADE rather than saturate: 1.5 rad/s produced +# -0.11 rad/s, i.e. the wrong direction, so the clamps below are real limits and +# not politeness. +GO1_MAX_FWD_MS = 1.0 +GO1_MAX_REV_MS = 0.4 +GO1_MAX_TURN_RADS = 0.8 +GO1_CTRL_DT = 0.02 # policy rate; 5 physics steps at our 4 ms timestep +GO1_ACTION_SCALE = 0.5 +GO1_FALLEN_HEIGHT_M = 0.16 +GO1_NOMINAL_HEIGHT_M = 0.27 +GO1_LEGS = ("FR", "FL", "RR", "RL") +GO1_SEGMENTS = ("hip", "thigh", "calf") + + +MOVE_ARRIVE_M = 0.5 + + +def _move_cmds( + dx_nwu: float, dy_nwu: float, yaw_nwu: float, yaw_hold_nwu: float +) -> tuple[float, float, float]: + """Body commands (vx, vy_ned_right, wz_ned_right) for a heading-held move. + + Pure math so it is unit-testable. Runs INSIDE the bridge at physics rate: + a first version lived in the DimOS module ticking on wall time, and at + ~40x realtime every command acted on 6-12 sim-seconds of stale state -- + the loop wandered a robot 140 m off a 5 m strafe. Control must live in + the same clock domain as the plant. + """ + fwd_err = dx_nwu * math.cos(yaw_nwu) + dy_nwu * math.sin(yaw_nwu) + left_err = -dx_nwu * math.sin(yaw_nwu) + dy_nwu * math.cos(yaw_nwu) + right_err = -left_err + vx = max(-0.4, min(0.7, 0.9 * fwd_err)) + vy = max(-0.3, min(0.3, 0.9 * right_err)) + # The policy's low-speed deadband makes tiny commands a stop, not a creep. + if 0.0 < abs(vx) < 0.3 and abs(fwd_err) > MOVE_ARRIVE_M / 2: + vx = math.copysign(0.3, vx) + if 0.0 < abs(vy) < 0.15 and abs(right_err) > MOVE_ARRIVE_M / 2: + vy = math.copysign(0.15, vy) + yaw_err = math.atan2(math.sin(yaw_hold_nwu - yaw_nwu), math.cos(yaw_hold_nwu - yaw_nwu)) + # NWU positive yaw is a LEFT turn; the NED wz convention is +right. + wz_ned = -max(-0.5, min(0.5, 1.5 * yaw_err)) + return vx, vy, wz_ned + + +class Go1Link: + """One real Unitree Go1, driven by the trained ONNX policy. + + Public surface is identical to :class:`LeggedLink` -- same UDP command + protocol, same state dict -- so ``LeggedSimModule`` and the swarm + coordinator do not know or care which one is underneath. + + The policy that ships with DimOS assumes it owns the whole model: it reads + ``qpos[7:]`` and writes ``ctrl[:]``. In this world that would read a drone's + joints and overwrite four drones' motor commands, so every lookup here is + resolved by id under the robot's own ``dogN-`` prefix and only that robot's + 12 actuators are ever written. + """ + + def __init__( + self, mj: Any, model: Any, data: Any, index: int, host: str, port: int, policy_path: str + ) -> None: + import numpy as np + import onnxruntime as ort + + self._np = np + self.index = index + self.name = f"dog{index}" + self._mj = mj + self.model = model + self.data = data + self.host = host + self.port = port + + prefix = f"dog{index}-" + n2i = mj.mj_name2id + self.body = n2i(model, mj.mjtObj.mjOBJ_BODY, prefix + "trunk") + if self.body < 0: + raise ValueError(f"scene has no {prefix}trunk body") + # The Go1's floating base joint is unnamed in the MJCF, so it has to be + # reached through the body rather than looked up by name. + free_j = int(model.body_jntadr[self.body]) + self.qpos_adr = int(model.jnt_qposadr[free_j]) + self.dof_adr = int(model.jnt_dofadr[free_j]) + + self._jq: list[int] = [] + self._jv: list[int] = [] + self._act: list[int] = [] + for leg in GO1_LEGS: + for seg in GO1_SEGMENTS: + j = n2i(model, mj.mjtObj.mjOBJ_JOINT, f"{prefix}{leg}_{seg}_joint") + a = n2i(model, mj.mjtObj.mjOBJ_ACTUATOR, f"{prefix}{leg}_{seg}") + if j < 0 or a < 0: + raise ValueError(f"scene missing {prefix}{leg}_{seg}") + self._jq.append(int(model.jnt_qposadr[j])) + self._jv.append(int(model.jnt_dofadr[j])) + self._act.append(a) + + self._sensor: dict[str, tuple[int, int]] = {} + for sname in ("local_linvel", "gyro"): + sid = n2i(model, mj.mjtObj.mjOBJ_SENSOR, prefix + sname) + if sid < 0: + raise ValueError(f"scene missing sensor {prefix}{sname}") + self._sensor[sname] = (int(model.sensor_adr[sid]), int(model.sensor_dim[sid])) + self._imu_site = n2i(model, mj.mjtObj.mjOBJ_SITE, prefix + "imu") + + self._default = np.array(GO1_HOME_ANGLES, dtype=np.float32) + self._session = ort.InferenceSession( + policy_path, providers=["CPUExecutionProvider"] + ) + self._last_action = np.zeros(12, dtype=np.float32) + self._since_ctrl = GO1_CTRL_DT # run the policy on the first step + + self.cmd_vx = 0.0 + self.cmd_vy = 0.0 + self.cmd_wz = 0.0 + self.fallen = False + self._diverged = False + self._policy_failed = False + # Heading-held relative move: (x_nwu, y_nwu, yaw_hold_nwu) or None. + self._move_target: tuple[float, float, float] | None = None + + # Stand at the home pose so the first observation is sane. + self.data.qpos[self.qpos_adr + 2] = GO1_SPAWN_Z + for k, adr in enumerate(self._jq): + self.data.qpos[adr] = float(self._default[k]) + for k, a in enumerate(self._act): + self.data.ctrl[a] = float(self._default[k]) + + self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self.sock.bind((host, port)) + self.sock.setblocking(False) + self.peer: tuple[str, int] | None = None + + # -- commands ---------------------------------------------------------- + + def poll_commands(self) -> None: + while True: + try: + payload, addr = self.sock.recvfrom(4096) + except (BlockingIOError, InterruptedError, OSError): + return + self.peer = addr + try: + cmd = json.loads(payload) + except ValueError: + logger.warning(f"[{self.name}] malformed command datagram") + continue + action = cmd.get("action", "") + if action == "recover": + self.fallen = False + self.cmd_vx = self.cmd_vy = self.cmd_wz = 0.0 + logger.info(f"[{self.name}] recovery attempt") + elif action == "stance": + # The shipped policy is a flat-ground velocity controller with no + # height input, so crouch/set_height have nothing to drive. Say + # so rather than silently accepting and doing nothing. + logger.info(f"[{self.name}] stance height not supported by the Go1 policy") + elif action == "move_to": + # Target arrives in NED; the bridge world is NWU (y = -east, + # yaw sign flips). + tn = float(cmd.get("north", 0.0)) + te = float(cmd.get("east", 0.0)) + yaw_ned = float(cmd.get("yaw", 0.0)) + self._move_target = (tn, -te, -yaw_ned) + logger.info(f"[{self.name}] move_to NED ({tn:.1f}, {te:.1f}), heading held") + elif action == "walk": + self.cmd_vx = max(-GO1_MAX_REV_MS, min(GO1_MAX_FWD_MS, float(cmd.get("vx", 0.0)))) + self._move_target = None # manual walk overrides an active move + # NED convention at the DimOS boundary: positive vy = strafe + # RIGHT. Clamped to the measured usable strafe (~0.3 m/s). + self.cmd_vy = max(-0.3, min(0.3, float(cmd.get("vy", 0.0)))) + self.cmd_wz = max( + -GO1_MAX_TURN_RADS, min(GO1_MAX_TURN_RADS, float(cmd.get("wz", 0.0))) + ) + elif action == "stop": + self._move_target = None + self.cmd_vx = self.cmd_vy = self.cmd_wz = 0.0 + elif action == "noop": + pass + + # -- policy ------------------------------------------------------------ + + def _observe(self) -> Any: + np = self._np + la, ld = self._sensor["local_linvel"] + ga, gd = self._sensor["gyro"] + linvel = self.data.sensordata[la : la + ld] + gyro = self.data.sensordata[ga : ga + gd] + rot = self.data.site_xmat[self._imu_site].reshape(3, 3) + gravity = rot.T @ np.array([0.0, 0.0, -1.0]) + angles = np.array([self.data.qpos[a] for a in self._jq]) - self._default + vels = np.array([self.data.qvel[a] for a in self._jv]) + # NED sign convention at the DimOS boundary: a positive commanded yaw + # rate means "nose right", which is negative about MuJoCo's z. + # cmd_vy and cmd_wz are NED-signed (positive = right); the policy was + # trained with +y = left and +yaw = counter-clockwise, so both negate. + command = np.array([self.cmd_vx, -self.cmd_vy, -self.cmd_wz], dtype=np.float32) + return np.hstack( + [linvel, gyro, gravity, angles, vels, self._last_action, command] + ).astype(np.float32) + + def apply(self, dt: float) -> None: + height = float(self.data.xpos[self.body][2]) + if height < GO1_FALLEN_HEIGHT_M: + if not self.fallen: + logger.warning(f"[{self.name}] fallen (h={height:.2f} m)") + self.fallen = True + self.cmd_vx = self.cmd_vy = self.cmd_wz = 0.0 + elif self.fallen and height > GO1_FALLEN_HEIGHT_M + 0.06: + logger.info(f"[{self.name}] upright again (h={height:.2f} m)") + self.fallen = False + + if self._move_target is not None and not self.fallen: + xt, yt, yaw_hold = self._move_target + pos = self.data.xpos[self.body] + dx, dy = xt - float(pos[0]), yt - float(pos[1]) + if math.hypot(dx, dy) <= MOVE_ARRIVE_M: + self._move_target = None + self.cmd_vx = self.cmd_vy = self.cmd_wz = 0.0 + logger.info(f"[{self.name}] move arrived ({math.hypot(dx, dy):.2f} m)") + else: + self.cmd_vx, self.cmd_vy, self.cmd_wz = _move_cmds( + dx, dy, self._yaw(), yaw_hold + ) + elif self._move_target is not None and self.fallen: + self._move_target = None + + self._since_ctrl += dt + if self._since_ctrl < GO1_CTRL_DT: + return + self._since_ctrl = 0.0 + + np = self._np + obs = self._observe() + # A non-finite observation means the physics for this robot has already + # diverged. Feeding it to the policy produces garbage torques that make + # the blow-up worse and can NaN the whole shared world -- including the + # drones. Hold the default stance instead and say so once. + if not np.all(np.isfinite(obs)): + if not self._diverged: + logger.error(f"[{self.name}] non-finite state; holding stance") + self._diverged = True + for k, a in enumerate(self._act): + self.data.ctrl[a] = float(self._default[k]) + return + self._diverged = False + + try: + self._last_action = self._session.run(None, {"obs": obs.reshape(1, -1)})[0][0] + except Exception as exc: + # One robot's policy failing must not take the world with it, for + # the same reason a dying autopilot does not (see _collect_replies). + if not self._policy_failed: + logger.error(f"[{self.name}] policy inference failed: {exc}; holding stance") + self._policy_failed = True + for k, a in enumerate(self._act): + self.data.ctrl[a] = float(self._default[k]) + return + self._policy_failed = False + + targets = self._last_action * GO1_ACTION_SCALE + self._default + for k, a in enumerate(self._act): + self.data.ctrl[a] = float(targets[k]) + + # -- telemetry --------------------------------------------------------- + + def _yaw(self) -> float: + qw, qx, qy, qz = self.data.qpos[self.qpos_adr + 3 : self.qpos_adr + 7] + return math.atan2(2.0 * (qw * qz + qx * qy), 1.0 - 2.0 * (qy * qy + qz * qz)) + + def state(self, sim_time_us: int) -> dict[str, Any]: + pos = self.data.xpos[self.body] + vel = self.data.qvel[self.dof_adr : self.dof_adr + 3] + return { + "key": self.name, + "robot_class": "legged", + "connected": True, + "nwu": [float(pos[0]), float(pos[1]), float(pos[2])], + "velocity": [float(vel[0]), float(vel[1]), float(vel[2])], + "yaw_rad": self._yaw(), + "height_m": float(pos[2]), + "nominal_height_m": GO1_NOMINAL_HEIGHT_M, + "walking": self._move_target is not None + or abs(self.cmd_vx) > 1e-3 + or abs(self.cmd_wz) > 1e-3 + or abs(self.cmd_vy) > 1e-3, + "fallen": self.fallen, + "speed_mps": self.cmd_vx, + "turn_rate_rads": self.cmd_wz, + "model": "unitree_go1", + "sim_time_s": sim_time_us / 1e6, + } + + def publish(self, sim_time_us: int) -> None: + if self.peer is None: + return + try: + self.sock.sendto(json.dumps(self.state(sim_time_us)).encode(), self.peer) + except OSError: + pass + + def close(self) -> None: + try: + self.sock.close() + except OSError: + pass + + +class Px4HilFleet: + """Drives N PX4 instances and M legged robots inside one MuJoCo world.""" + + def __init__( + self, + n_drones: int = 3, + n_dogs: int = 1, + host: str = "127.0.0.1", + port_base: int = 4560, + legged_port_base: int = LEGGED_PORT_BASE, + origin: tuple[float, float, float] = ( + DEFAULT_ORIGIN_LAT, + DEFAULT_ORIGIN_LON, + DEFAULT_ORIGIN_ALT, + ), + viewer: bool = False, + ) -> None: + import mujoco + + self._mj = mujoco + self.origin = origin + self.want_viewer = viewer + + self.model, self._real_go1 = build_model(n_drones, n_dogs) + self._policy_path = go1_policy_path() if self._real_go1 else None + if self._real_go1 and self._policy_path is None: + # Real mesh but no trained weights would silently fall back to a + # robot that cannot walk at all, so rebuild the primitive instead. + logger.warning("Go1 mesh found but no trained policy; using primitive quadruped") + self.model, self._real_go1 = build_model(n_drones, n_dogs, real_go1=False) + self.data = mujoco.MjData(self.model) + # Populate sensordata and site frames before anything is transmitted. + # PX4 aligns its initial attitude from the first sample it receives, and + # an all-zero one leaves the EKF stably but completely wrong. + mujoco.mj_forward(self.model, self.data) + + if n_drones and not x500_meshes_available(): + logger.info( + f"drones: primitive airframe x{n_drones} (flies identically); " + "run tools/fetch_x500_meshes.py for the X500 model" + ) + self.links = [ + DroneLink(mujoco, self.model, self.data, i, host, port_base + i) + for i in range(n_drones) + ] + if self._real_go1 and self._policy_path is not None: + logger.info(f"legged robots: real Unitree Go1 + trained policy x{n_dogs}") + self.dogs: list[Any] = [ + Go1Link( + mujoco, self.model, self.data, i, host, + legged_port_base + i, self._policy_path, + ) + for i in range(n_dogs) + ] + else: + if n_dogs: + # Say WHY. A silent downgrade leaves a robot that walks badly + # and no clue that a far better one was one command away. + reason = go1_unavailable_reason() or "trained policy missing" + logger.warning( + f"legged robots: PRIMITIVE quadruped x{n_dogs} " + f"(0.37 m/s, falls on reversals) -- real Go1 unavailable: {reason}" + ) + self.dogs = [ + LeggedLink(mujoco, self.model, self.data, i, host, legged_port_base + i) + for i in range(n_dogs) + ] + # The Go1 links write their home stance into qpos during construction, + # so the world has to be re-evaluated before the first sensor read. + mujoco.mj_forward(self.model, self.data) + for dog in self.dogs: + dog.poll_commands() + dog.apply(self.model.opt.timestep) + + # One wind for the shared world: constant mean (SIM_WIND_N/E, m/s in + # the NED sense: N = from where it blows TO... no -- the vector the + # air MOVES with, world frame: +N pushes everything north) plus an + # Ornstein-Uhlenbeck gust (SIM_GUST_STD, m/s). Zero by default. + import os as _os + + self._wind_mean = np.array( + [ + float(_os.getenv("SIM_WIND_N", "0") or 0), + -float(_os.getenv("SIM_WIND_E", "0") or 0), # world is NWU: y = -E + 0.0, + ] + ) + self._gust_std = float(_os.getenv("SIM_GUST_STD", "0") or 0) + self._gust = np.zeros(3) + self._wind_rng = np.random.default_rng(7) + if np.any(self._wind_mean) or self._gust_std: + logger.info( + f"wind enabled: mean N={self._wind_mean[0]:.1f} E={-self._wind_mean[1]:.1f} m/s, " + f"gust std {self._gust_std:.1f} m/s" + ) + + self.sim_time_us = 1_000_000 + self.step_us = round(self.model.opt.timestep * 1e6) + self.gps_interval_us = int(1e6 / GPS_RATE_HZ) + self._last_gps_us = 0 + self._legged_interval_us = int(1e6 / LEGGED_STATE_HZ) + self._last_legged_us = 0 + + # -- lifecycle --------------------------------------------------------- + + def wait_for_px4(self, timeout_s: float = 180.0) -> None: + for link in self.links: + link.listen() + if self.links: + ports = ", ".join(str(link.port) for link in self.links) + logger.info(f"listening on {ports}; waiting for {len(self.links)} PX4 instance(s)") + if self.dogs: + logger.info( + "legged endpoints: " + + ", ".join(f"{d.name}=udp:{d.host}:{d.port}" for d in self.dogs) + ) + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if all(link.try_accept() for link in self.links): + logger.info(f"all {len(self.links)} PX4 instance(s) connected") + return + missing = [link.name for link in self.links if link.conn is None] + raise TimeoutError(f"PX4 never connected for: {', '.join(missing)}") + + GROUND_PROX_HALT_M = 0.6 + GROUND_PROX_REARM_M = 1.2 + + def _ground_proximity_guard(self) -> None: + """Halt any two ground robots about to walk into each other. + + The aircraft get dispatch-time separation; ground robots deliberately + have no separation floor (slow, low stakes) -- but "two gotos to the + same point" should still not end in robots pushing at each other. This + runs in the bridge at sim rate, so it also catches what no dispatch + check can: two already-moving robots converging. Hysteresis: halted + pairs re-arm once they are GROUND_PROX_REARM_M apart, so the operator + can walk one away without the guard re-firing every step. + """ + if len(self.dogs) < 2: + return + halted: set[tuple[int, int]] = getattr(self, "_prox_halted", set()) + for i in range(len(self.dogs)): + for j in range(i + 1, len(self.dogs)): + a, b = self.dogs[i], self.dogs[j] + pa, pb = self.data.xpos[a.body], self.data.xpos[b.body] + d = math.hypot(float(pa[0] - pb[0]), float(pa[1] - pb[1])) + key = (i, j) + moving = any( + abs(getattr(lk, "cmd_vx", 0.0)) > 1e-3 + or abs(getattr(lk, "cmd_vy", 0.0)) > 1e-3 + or getattr(lk, "_move_target", None) is not None + for lk in (a, b) + ) + if key not in halted and d < self.GROUND_PROX_HALT_M and moving: + for lk in (a, b): + lk.cmd_vx = lk.cmd_wz = 0.0 + if hasattr(lk, "cmd_vy"): + lk.cmd_vy = 0.0 + if getattr(lk, "_move_target", None) is not None: + lk._move_target = None + halted.add(key) + logger.warning( + f"GROUND PROXIMITY: {a.name} and {b.name} {d:.2f} m apart " + "and converging -- both halted" + ) + elif key in halted and d > self.GROUND_PROX_REARM_M: + halted.discard(key) + logger.info(f"ground proximity cleared: {a.name}/{b.name} ({d:.2f} m)") + self._prox_halted = halted + + def _collect_replies(self) -> None: + """Wait for every live vehicle's actuator reply, all sockets at once. + + Polling each vehicle with its own blocking recv costs a scheduler + round-trip per vehicle per step and was the dominant cost in the loop. + """ + pending = [lk for lk in self.links if lk.conn is not None] + for lk in pending: + lk.got_reply = False + deadline = time.monotonic() + ACTUATOR_WAIT_S + while pending: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + readable, _, _ = select.select([lk.conn for lk in pending], [], [], remaining) + except (OSError, ValueError): + break + if not readable: + break + for lk in list(pending): + if lk.conn in readable: + try: + if lk.pump(): + pending.remove(lk) + except ConnectionError as e: + # One autopilot dying must not take the world with it. + logger.warning(f"{e}; continuing without it") + lk.close() + pending.remove(lk) + + def step(self) -> None: + send_gps = self.sim_time_us - self._last_gps_us >= self.gps_interval_us + for link in self.links: + try: + link.send_sensors(self.sim_time_us, self.origin) + if send_gps: + link.send_gps(self.sim_time_us, self.origin) + except ConnectionError as e: + logger.warning(f"{e}; continuing without it") + link.close() + if send_gps: + self._last_gps_us = self.sim_time_us + + self._collect_replies() + if self._gust_std: + from dimos.simulation.px4_hil.hil_bridge import GUST_TAU_S + + dt = float(self.model.opt.timestep) + self._gust[:2] += (-self._gust[:2] / GUST_TAU_S) * dt + self._gust_std * math.sqrt( + 2.0 * dt / GUST_TAU_S + ) * self._wind_rng.standard_normal(2) + wind = self._wind_mean + self._gust + for link in self.links: + link.apply(wind) + + # Legged robots: take any waiting command, then advance the gait by one + # physics step. Both must happen every step -- polling only at startup + # means the bridge never learns a module's address and never publishes + # state back to it, and skipping apply() freezes the gait mid-swing. + for dog in self.dogs: + dog.poll_commands() + dog.apply(self.model.opt.timestep) + + self._ground_proximity_guard() + + if self.sim_time_us - self._last_legged_us >= self._legged_interval_us: + for dog in self.dogs: + dog.publish(self.sim_time_us) + self._last_legged_us = self.sim_time_us + + self._mj.mj_step(self.model, self.data) + self.sim_time_us += self.step_us + + def run(self) -> None: + self.wait_for_px4() + viewer_ctx = None + try: + if self.want_viewer: + import mujoco.viewer + + viewer_ctx = mujoco.viewer.launch_passive(self.model, self.data) + last = time.monotonic() + steps = 0 + while True: + self.step() + steps += 1 + if viewer_ctx is not None and steps % 25 == 0: + if not viewer_ctx.is_running(): + # Closing the window must not take the fleet with it. + # Before this check, an accidental click on the X killed + # the physics while PX4 and the daemon kept running, + # leaving a fleet of ghosts that answered but never + # moved. Drop to headless instead. + viewer_ctx.close() + viewer_ctx = None + logger.info("viewer closed; continuing headless (sim.sh stop to shut down)") + else: + viewer_ctx.sync() + now = time.monotonic() + if now - last >= 5.0: + rtf = (steps * self.model.opt.timestep) / (now - last) + logger.info( + f"sim={self.sim_time_us / 1e6:7.1f}s rtf={rtf:5.2f}x " + + " ".join( + f"{lk.name}(alt={self.data.xpos[lk.body][2]:5.2f}," + f"armed={int(lk.armed)})" + for lk in self.links + ) + + "".join( + f" {d.name}(z={self.data.xpos[d.body][2]:4.2f})" for d in self.dogs + ) + ) + last, steps = now, 0 + except (KeyboardInterrupt, ConnectionError, TimeoutError) as e: + logger.info(f"fleet bridge stopping: {e or 'interrupted'}") + finally: + if viewer_ctx is not None: + viewer_ctx.close() + for link in self.links: + link.close() + for dog in self.dogs: + dog.close() + + +def main() -> None: + p = argparse.ArgumentParser(description="PX4 SITL fleet + legged robots <-> MuJoCo") + p.add_argument("--drones", type=int, default=3) + p.add_argument("--dogs", type=int, default=1) + p.add_argument("--host", default="127.0.0.1") + p.add_argument("--port-base", type=int, default=4560) + p.add_argument("--legged-port-base", type=int, default=LEGGED_PORT_BASE) + p.add_argument("--lat", type=float, default=DEFAULT_ORIGIN_LAT) + p.add_argument("--lon", type=float, default=DEFAULT_ORIGIN_LON) + p.add_argument("--alt", type=float, default=DEFAULT_ORIGIN_ALT) + p.add_argument("--viewer", action="store_true") + a = p.parse_args() + Px4HilFleet( + n_drones=a.drones, n_dogs=a.dogs, host=a.host, port_base=a.port_base, + legged_port_base=a.legged_port_base, origin=(a.lat, a.lon, a.alt), viewer=a.viewer, + ).run() + + +if __name__ == "__main__": + main() diff --git a/dimos/simulation/px4_hil/hil_bridge.py b/dimos/simulation/px4_hil/hil_bridge.py new file mode 100644 index 0000000000..e19122a3fd --- /dev/null +++ b/dimos/simulation/px4_hil/hil_bridge.py @@ -0,0 +1,432 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MAVLink HIL protocol: constants, frame conversions, and one vehicle's link. + +Gazebo talks to PX4 through ``gz_bridge``, a module compiled into the firmware. +MuJoCo has no equivalent, so this implements the other supported path: PX4's +generic HIL interface (``simulator_mavlink``), the same one jMAVSim, JSBSim and +FlightGear use. The firmware is unmodified, which is the property the hardware +deliverable depends on -- what flies here is what flies on a Pixhawk. + +Protocol, from PX4's SimulatorMavlink.cpp: + +* Transport is **TCP on 4560 + instance**, and *PX4 is the client*. It sits in a + connect() retry loop at startup, so we must be listening first. +* Sim -> PX4: ``HIL_SENSOR`` every step, ``HIL_GPS`` at a lower rate. +* PX4 -> Sim: ``HIL_ACTUATOR_CONTROLS``, normalised motor outputs, zero while + disarmed. + +Lockstep is the part worth understanding. On every ``HIL_SENSOR`` with +``id == 0`` PX4 calls ``px4_clock_settime()`` with our ``time_usec`` and then +``px4_lockstep_progress()``. **The simulator owns PX4's clock.** PX4 never runs +on wall-clock time; it advances exactly as fast as we feed it. A slow host makes +the simulation run slow rather than starving the sensor stream -- the failure +that caps the Gazebo path at one flyable drone on a loaded machine. + +Frames, which is where this kind of bridge usually goes wrong: + +* MuJoCo body frame is **FLU** (x-Forward, y-Left, z-Up); PX4 body is **FRD**. +* MuJoCo world frame is **NWU** (x-North, y-West, z-Up); PX4 world is **NED**. + +Both conversions are ``(x, -y, -z)``. NWU rather than the more usual ENU is +deliberate: PX4 assumes a body at zero yaw points North, and a body at identity +in MuJoCo points along world +x. An ENU world puts "forward" at East, which +yaws the magnetometer 90 degrees and diverges the attitude estimate. + +The runnable entry point is :mod:`dimos.simulation.px4_hil.fleet_bridge`, which +drives any number of these links (including exactly one) against one shared +MuJoCo world. +""" + +from __future__ import annotations + +import math +import socket +from typing import Any + +import numpy as np + +from dimos.simulation.px4_hil.scene import MAX_THRUST_PER_ROTOR_N + +# --- HIL_SENSOR.fields_updated bitmask ------------------------------------- +# Verbatim from SimulatorMavlink.hpp `enum class SensorSource`. PX4 ignores any +# sensor group whose bits are not set, so getting these wrong shows up as "no +# valid data from Baro 0" rather than as an obvious protocol error. +FIELD_ACCEL = 0b111 # xacc, yacc, zacc +FIELD_GYRO = 0b111000 # xgyro, ygyro, zgyro +FIELD_MAG = 0b111000000 # xmag, ymag, zmag +FIELD_BARO = 0b1101000000000 # abs_pressure, pressure_alt, temperature +FIELDS_MULTIROTOR = FIELD_ACCEL | FIELD_GYRO | FIELD_MAG | FIELD_BARO + +# PX4's default SITL home (Zurich Irchel), so QGroundControl and the Gazebo-era +# blueprints agree on where "home" is. +DEFAULT_ORIGIN_LAT = 47.397742 +DEFAULT_ORIGIN_LON = 8.545594 +DEFAULT_ORIGIN_ALT = 488.0 + +# Earth magnetic field at that origin, in NED, Gauss. A wrong magnitude or sign +# shows up as the EKF refusing to converge on yaw, which reads like a control +# bug rather than a sensor one. +MAG_FIELD_NED_GAUSS = (0.21523, 0.0, 0.42980) + +EARTH_RADIUS_M = 6371000.0 +SEA_LEVEL_PRESSURE_HPA = 1013.25 +SEA_LEVEL_TEMP_C = 15.0 +TEMP_LAPSE_RATE_C_PER_M = 0.0065 + +# --- Sensor noise ---------------------------------------------------------- +# EKF2 is tuned for real, noisy sensors and *needs* that noise. Fed a perfectly +# deterministic stream its covariance collapses and any small model mismatch +# lands in the bias states instead: an invented accelerometer bias, a yaw that +# settles far from truth, a position that drifts kilometres at zero velocity. +# gz_bridge and jMAVSim inject noise for exactly this reason. Values are +# 1-sigma per sample at 250 Hz (sigma = density * sqrt(rate)). +NOISE_ACCEL_MS2 = 0.03 # 1.86e-3 m/s^2/sqrt(Hz) +NOISE_GYRO_RADS = 0.003 # 1.87e-4 rad/s/sqrt(Hz) +NOISE_MAG_GAUSS = 0.0004 +NOISE_BARO_HPA = 0.01 +NOISE_GPS_POS_M = 0.03 # well inside the eph/epv we advertise +NOISE_GPS_VEL_MS = 0.02 + +# HIL_GPS is expensive relative to HIL_SENSOR, and real receivers are slow. +GPS_RATE_HZ = 10.0 + +# Deadline for collecting every vehicle's actuator reply in one lockstep cycle. +# PX4 publishes nothing until its output modules are up, so early cycles time +# out legitimately; this only has to be long enough to cover a scheduler +# round-trip once things are running. +ACTUATOR_WAIT_S = 0.05 + + +def _flu_to_frd(v: Any) -> tuple[float, float, float]: + """MuJoCo body frame (Forward-Left-Up) -> PX4 body frame (Forward-Right-Down).""" + return (float(v[0]), float(-v[1]), float(-v[2])) + + +def _nwu_to_ned(v: Any) -> tuple[float, float, float]: + """MuJoCo world frame (North-West-Up) -> PX4 world frame (North-East-Down).""" + return (float(v[0]), float(-v[1]), float(-v[2])) + + +def _pressure_hpa(altitude_m: float) -> float: + """Barometric pressure at an altitude above mean sea level, hPa.""" + base_k = SEA_LEVEL_TEMP_C + 273.15 + temp_k = base_k - TEMP_LAPSE_RATE_C_PER_M * altitude_m + return SEA_LEVEL_PRESSURE_HPA * (temp_k / base_k) ** 5.25588 + + +# --------------------------------------------------------------------------- +# Rotor-craft aerodynamics -- the "necessary physics" layer +# --------------------------------------------------------------------------- +# MuJoCo's built-in density-based drag covers the airframe as a bluff body, but +# none of what makes a multirotor feel like one. These are the standard +# first-order additions every serious PX4 sim carries, with the standard +# approximations, all computed in SIM time inside the bridge: +# +# * motor lag -- ESC+prop spool is a first-order response, not a step. +# * rotor H-drag -- a translating rotor disc produces an in-plane force +# opposing airspeed, proportional to thrust. THE dominant +# damping on a multirotor; without it, braking and wind +# response are wrong. +# * ground effect -- thrust rises near the ground (classic image model), +# felt as float in the last half-metre of a landing. +# * wind + gusts -- constant wind (SIM_WIND_N/SIM_WIND_E, m/s) plus an +# Ornstein-Uhlenbeck gust process (SIM_GUST_STD, m/s). +# Enters through the airspeed the H-drag sees. +MOTOR_TAU_S = 0.06 # ESC+prop spool time constant (typ. 0.02-0.1) +ROTOR_HDRAG_S_PER_M = 0.06 # H-force coefficient (typ. 0.03-0.10) +ROTOR_RADIUS_M = 0.127 # 10-inch prop +GROUND_EFFECT_MAX = 0.12 # cap the boost at 12% of current thrust +GUST_TAU_S = 2.0 # gust correlation time + + +def _motor_lag_step(state: np.ndarray, command: np.ndarray, dt: float) -> np.ndarray: + """First-order spool toward `command`; returns the new state.""" + alpha = 1.0 - math.exp(-dt / MOTOR_TAU_S) + return state + alpha * (command - state) + + +def _ground_effect_boost(height_m: float) -> float: + """Extra thrust fraction from the classic image-rotor model, capped. + + boost = 1/(1 - (R/4h)^2) - 1, clamped below h = R/2 where the model blows + up (a landed vehicle is not hovering in its own wake). + """ + h = max(height_m, ROTOR_RADIUS_M / 2.0) + ratio = ROTOR_RADIUS_M / (4.0 * h) + return min(GROUND_EFFECT_MAX, 1.0 / (1.0 - ratio * ratio) - 1.0) + + +def _rotor_hdrag_force(v_air_world: np.ndarray, thrust_n: float) -> np.ndarray: + """In-plane rotor drag, world frame: -c * T * v_air (xy only). + + Proportional to thrust because the H-force scales with blade lift; zero + thrust (motors off, falling) produces none, which is also correct. + """ + f = -ROTOR_HDRAG_S_PER_M * thrust_n * v_air_world + f[2] = 0.0 + return f + + +class DroneLink: + """One PX4 instance's HIL endpoint, bound to one body in a shared world. + + Owns the socket and the MuJoCo handles for a single vehicle; the caller owns + the world and the stepping. Keeping the loop out of here is what lets one + process drive an arbitrary number of vehicles in lockstep together. + """ + + def __init__(self, mj: Any, model: Any, data: Any, index: int, host: str, port: int) -> None: + self.index = index + self.name = f"drone{index}" + self.host = host + self.port = port + self._mj = mj + self.model = model + self.data = data + + n2i = mj.mj_name2id + self.site = n2i(model, mj.mjtObj.mjOBJ_SITE, f"{self.name}_imu") + self.body = n2i(model, mj.mjtObj.mjOBJ_BODY, f"{self.name}_base") + if self.site < 0 or self.body < 0: + raise ValueError(f"scene has no body/site for {self.name}") + self.acc_adr = int(model.sensor_adr[n2i(model, mj.mjtObj.mjOBJ_SENSOR, f"{self.name}_acc")]) + self.gyro_adr = int( + model.sensor_adr[n2i(model, mj.mjtObj.mjOBJ_SENSOR, f"{self.name}_gyro")] + ) + jid = n2i(model, mj.mjtObj.mjOBJ_JOINT, f"{self.name}_root") + self.dof_adr = int(model.jnt_dofadr[jid]) + self.motor_ids = [ + n2i(model, mj.mjtObj.mjOBJ_ACTUATOR, f"{self.name}_motor{k}") for k in range(4) + ] + + self.sock: socket.socket | None = None + self.conn: socket.socket | None = None + self.mav: Any = None + self.controls = np.zeros(4) + # Filtered motor state (what the "motors" actually produce after lag). + self.motor_state = np.zeros(4) + # Visual handles are optional -- absent on the primitive airframe. + gid = lambda n: n2i(model, mj.mjtObj.mjOBJ_GEOM, n) # noqa: E731 + self.prop_gids = [gid(f"{self.name}_prop{k}") for k in range(4)] + self.disc_gids = [gid(f"{self.name}_disc{k}") for k in range(4)] + self.led_gid = gid(f"{self.name}_led") + self._have_fx = all(g >= 0 for g in self.prop_gids + self.disc_gids) + self._prop_angle = 0.0 + self.armed = False + self.rx_actuator = 0 + self.got_reply = False + self.rng = np.random.default_rng(1000 + index) + # Scratch buffers, reused every step. At 250 Hz x N vehicles the + # allocation churn from fresh noise arrays is measurable. + self._noise3 = np.empty(3) + + # -- connection -------------------------------------------------------- + + def listen(self) -> None: + self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self.sock.bind((self.host, self.port)) + self.sock.listen(1) + self.sock.settimeout(0.5) + + def try_accept(self) -> bool: + if self.conn is not None: + return True + assert self.sock is not None + try: + conn, _addr = self.sock.accept() + except (TimeoutError, BlockingIOError): + return False + from pymavlink.dialects.v20 import common as mavlink + + # PX4 sets TCP_NODELAY on its side; match it or lockstep pays a Nagle + # delay on every single exchange. + conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + conn.setblocking(False) + self.conn = conn + self.mav = mavlink.MAVLink(None, srcSystem=1, srcComponent=1) + return True + + def fileno(self) -> int: + return self.conn.fileno() if self.conn is not None else -1 + + def close(self) -> None: + for s in (self.conn, self.sock): + if s is not None: + try: + s.close() + except OSError: + pass + self.conn = None + + # -- sim -> PX4 -------------------------------------------------------- + + def send_sensors(self, sim_time_us: int, origin: tuple[float, float, float]) -> None: + from pymavlink.dialects.v20 import common as mavlink + + if self.conn is None: + return + d = self.data + rng = self.rng + acc = d.sensordata[self.acc_adr : self.acc_adr + 3] + gyro = d.sensordata[self.gyro_adr : self.gyro_adr + 3] + xacc, yacc, zacc = _flu_to_frd(acc + rng.normal(0, NOISE_ACCEL_MS2, 3)) + xgyro, ygyro, zgyro = _flu_to_frd(gyro + rng.normal(0, NOISE_GYRO_RADS, 3)) + + # site_xmat is body->world (NWU); its transpose takes world->body. + r_body_to_world = d.site_xmat[self.site].reshape(3, 3) + n, e, dn = MAG_FIELD_NED_GAUSS + mag = r_body_to_world.T @ np.array([n, -e, -dn]) + rng.normal(0, NOISE_MAG_GAUSS, 3) + xmag, ymag, zmag = _flu_to_frd(mag) + + altitude = float(d.xpos[self.body][2]) + origin[2] + self._send( + mavlink.MAVLink_hil_sensor_message( + time_usec=sim_time_us, + xacc=xacc, yacc=yacc, zacc=zacc, + xgyro=xgyro, ygyro=ygyro, zgyro=zgyro, + xmag=xmag, ymag=ymag, zmag=zmag, + abs_pressure=_pressure_hpa(altitude) + rng.normal(0, NOISE_BARO_HPA), + diff_pressure=0.0, + pressure_alt=altitude, + temperature=SEA_LEVEL_TEMP_C - TEMP_LAPSE_RATE_C_PER_M * altitude, + fields_updated=FIELDS_MULTIROTOR, + id=0, # primary IMU: this is the message that drives lockstep + ) + ) + + def send_gps(self, sim_time_us: int, origin: tuple[float, float, float]) -> None: + from pymavlink.dialects.v20 import common as mavlink + + if self.conn is None: + return + olat, olon, oalt = origin + d = self.data + pos = d.xpos[self.body] + self.rng.normal(0, NOISE_GPS_POS_M, 3) + vel = d.qvel[self.dof_adr : self.dof_adr + 3] + self.rng.normal(0, NOISE_GPS_VEL_MS, 3) + n, e, down = _nwu_to_ned(pos) + vn, ve, vd = _nwu_to_ned(vel) + lat = olat + math.degrees(n / EARTH_RADIUS_M) + lon = olon + math.degrees(e / (EARTH_RADIUS_M * math.cos(math.radians(olat)))) + self._send( + mavlink.MAVLink_hil_gps_message( + time_usec=sim_time_us, + fix_type=3, + lat=int(lat * 1e7), lon=int(lon * 1e7), alt=int((oalt - down) * 1e3), + eph=30, epv=40, + vel=int(math.hypot(vn, ve) * 100), + vn=int(vn * 100), ve=int(ve * 100), vd=int(vd * 100), + cog=int((math.degrees(math.atan2(ve, vn)) % 360.0) * 100), + satellites_visible=14, id=0, + ) + ) + + def _send(self, msg: Any) -> None: + assert self.conn is not None + try: + self.conn.sendall(msg.pack(self.mav)) + except (BlockingIOError, InterruptedError): + # Kernel send buffer full: PX4 is behind. Dropping this sample is + # better than stalling every other vehicle in the world. + pass + except OSError as e: + raise ConnectionError(f"[{self.name}] send failed: {e}") from e + + # -- PX4 -> sim -------------------------------------------------------- + + def pump(self) -> bool: + """Drain whatever is readable. Returns True if an actuator command arrived.""" + if self.conn is None or self.mav is None: + return False + got = False + while True: + try: + chunk = self.conn.recv(8192) + except (BlockingIOError, InterruptedError): + break + except OSError as e: + raise ConnectionError(f"[{self.name}] recv failed: {e}") from e + if not chunk: + raise ConnectionError(f"[{self.name}] PX4 closed the connection") + for msg in self.mav.parse_buffer(chunk) or []: + if msg.get_type() == "HIL_ACTUATOR_CONTROLS": + self.controls[:] = msg.controls[:4] + # bit 7 of `mode` is MAV_MODE_FLAG_SAFETY_ARMED. + self.armed = bool(msg.mode & 0b10000000) + self.rx_actuator += 1 + got = True + if len(chunk) < 8192: + break + self.got_reply = self.got_reply or got + return got + + def apply(self, wind_world: np.ndarray | None = None) -> None: + np.clip(self.controls, 0.0, 1.0, out=self.controls) + dt = float(self.model.opt.timestep) + # Motor spool: PX4 commands a step; the airframe answers with a lag. + self.motor_state = _motor_lag_step(self.motor_state, self.controls, dt) + for k, aid in enumerate(self.motor_ids): + self.data.ctrl[aid] = self.motor_state[k] + + thrust_n = float(self.motor_state.sum()) * MAX_THRUST_PER_ROTOR_N + v_world = np.array(self.data.qvel[self.dof_adr : self.dof_adr + 3]) + v_air = v_world if wind_world is None else v_world - wind_world + force = _rotor_hdrag_force(v_air, thrust_n) + height = float(self.data.xpos[self.body][2]) + force[2] += _ground_effect_boost(height) * thrust_n + self.data.xfrc_applied[self.body, :3] = force + + self._update_fx(dt) + + def _update_fx(self, dt: float) -> None: + """Drive the cosmetic layer: prop/disc cross-fade, spin, status LED. + + Pure rendering -- every geom touched is contype 0 / mass 0. Writes go + to the model's rgba/quat arrays, which the viewer re-reads each frame. + """ + if not self._have_fx: + return + throttle = float(self.motor_state.mean()) + # Cross-fade: still blades at rest, translucent disc at speed. + prop_a = max(0.15, 1.0 - 3.0 * throttle) + disc_a = min(0.38, 1.6 * throttle) + # Slow visible churn during spool-up sells the transition; the quat + # write is on the mesh geom only. + self._prop_angle = (self._prop_angle + 60.0 * throttle * dt) % (2.0 * math.pi) + half = self._prop_angle / 2.0 + cw, sw = math.cos(half), math.sin(half) + for k in range(4): + self.model.geom_rgba[self.prop_gids[k], 3] = prop_a + self.model.geom_rgba[self.disc_gids[k], 3] = disc_a + q = self.model.geom_quat[self.prop_gids[k]] + q[0], q[1], q[2], q[3] = cw, 0.0, 0.0, (sw if k < 2 else -sw) + if self.led_gid >= 0: + self.model.geom_rgba[self.led_gid, :3] = ( + (0.95, 0.15, 0.1) if self.armed else (0.1, 0.9, 0.2) + ) + + +__all__ = [ + "ACTUATOR_WAIT_S", + "DEFAULT_ORIGIN_ALT", + "DEFAULT_ORIGIN_LAT", + "DEFAULT_ORIGIN_LON", + "FIELDS_MULTIROTOR", + "GPS_RATE_HZ", + "MAG_FIELD_NED_GAUSS", + "DroneLink", +] diff --git a/dimos/simulation/px4_hil/scene.py b/dimos/simulation/px4_hil/scene.py new file mode 100644 index 0000000000..953ea1c52e --- /dev/null +++ b/dimos/simulation/px4_hil/scene.py @@ -0,0 +1,583 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Builds a mixed MuJoCo scene: N PX4 quadrotors plus M legged robots. + +The whole point of moving off Gazebo is that aircraft and ground robots end up +in *one* physics world, so they can see and avoid each other for real rather +than through a co-simulation bridge that only exchanges positions. + +Everything here is generated from primitives -- boxes, capsules, cylinders -- +with no mesh dependencies. DimOS's usual legged assets (Go1/G1) live in a +git-lfs blob and need `mujoco_playground`, neither of which is available on +every machine; a procedural quadruped keeps the scene runnable anywhere. The +real meshes can be swapped in later without touching the bridge, because the +bridge only ever refers to bodies and sensors by name. + +Naming is the contract between this file and the bridge. Every vehicle gets a +unique prefix, and the bridge resolves handles from it: + + drone0_imu, drone0_thrust0..3, drone0_acc, drone0_gyro (per quadrotor) + dog0_torso, dog0_imu (per quadruped) + +That prefix is deliberately the same string used for the DimOS namespace, so a +vehicle's identity is consistent from MuJoCo through MAVLink to the agent's +tool names. + +Frames: world is NWU (x-North, y-West, z-Up) to match hil_bridge.py, which +needs a body at identity to point North the way PX4 expects. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +# Rotor arm positions for a real Holybro X500: a symmetric 0.174 m square, +# taken from the x500_base SDF the visual mesh comes from. +# PX4 states them in FRD (y right); MuJoCo is FLU, so y is negated here. +# Order matters: index i is PX4 motor i+1, i.e. HIL_ACTUATOR_CONTROLS[i]. +# +# These MUST stay in lockstep with the CA_ROTOR*_P* values written by +# sim_params.py. PX4 builds its control-allocation matrix from those; if +# MuJoCo applies the thrust anywhere else the controller mis-allocates torque. +# (PX4's own 4001_gz_x500 ships (0.13, 0.22), which matches neither its own mesh +# nor the real airframe, so we override it on both sides instead of adopting it.) +X500_ARM_M = 0.174 + +# Rotor plane height above OUR body origin. +# +# The SDF puts the rotor links at z=0.06 above base_link, with the frame visual +# at z=+0.025. But base_link is not our body origin: we centre the frame mesh on +# its bounding box so the model rests on its landing gear, and that box centre +# sits 0.1131 m below the mesh origin (the gear hangs a long way down). So +# everything measured in the SDF frame shifts up by (0.1131 - 0.025) here. +# +# Getting this wrong is what left the props hanging 8.8 cm under the motors. +# Verified against the mesh: structure under each rotor position tops out at +# z=+0.025 in the SDF frame, i.e. 0.1131 here, and the props sit just above it. +X500_FRAME_BBOX_CZ = -0.1131 # bbox centre of the raw frame mesh +X500_SDF_VISUAL_DZ = 0.025 # SDF pose of the frame visual in base_link +_SDF_TO_BODY_Z = -(X500_FRAME_BBOX_CZ + X500_SDF_VISUAL_DZ) # = +0.0881 +X500_ROTOR_Z = 0.06 + _SDF_TO_BODY_Z # 0.1481 +X500_MOTOR_Z = X500_ROTOR_Z - 0.032 # SDF motor-base offset in the rotor link +ROTORS = [ + (X500_ARM_M, -X500_ARM_M, "0.8 0.2 0.2 1"), # M1 front-right + (-X500_ARM_M, X500_ARM_M, "0.2 0.2 0.8 1"), # M2 rear-left + (X500_ARM_M, X500_ARM_M, "0.2 0.8 0.2 1"), # M3 front-left + (-X500_ARM_M, -X500_ARM_M, "0.8 0.8 0.2 1"), # M4 rear-right +] +# Rotor i spins CCW when its PX4 KM is negative. M1/M2 share one sense, M3/M4 +# the other, so the CW/CCW prop meshes are assigned to match. +ROTOR_PROP_MESH = ("prop_cw", "prop_cw", "prop_ccw", "prop_ccw") +# +KM rotors spin one way, -KM the other. PX4 builds its effectiveness matrix as +# moment = ct * position.cross(axis) - ct * km * axis (axis = (0,0,-1) FRD) +# so torque_z(FRD) = +km * thrust. FRD z points down, so a positive FRD yaw +# torque is NEGATIVE about MuJoCo's z. Getting this backwards makes yaw a +# positive-feedback loop and the vehicle flips within half a second of arming. +ROTOR_YAW_TORQUE = [-0.368, -0.368, 0.368, 0.368] + +# Per-vehicle accent colours. Purely cosmetic, but load-bearing for a human +# watching the viewer: with every airframe the same grey you cannot tell which +# drone is executing which lane of a sweep. Index = vehicle number. +ACCENTS = [ + "0.95 0.26 0.21 1", # red + "0.16 0.71 0.96 1", # cyan + "1.00 0.76 0.03 1", # amber + "0.30 0.89 0.44 1", # green + "0.85 0.36 0.95 1", # violet + "1.00 0.50 0.15 1", # orange +] + + +def _accent(i: int) -> str: + return ACCENTS[i % len(ACCENTS)] +MAX_THRUST_PER_ROTOR_N = 7.36 + +# Quadruped standing pose. Slightly bent so the legs are not in the singular +# fully-extended configuration, which makes the stance solver ill-conditioned. +DOG_HIP_ANGLE = 0.3 +DOG_KNEE_ANGLE = -0.6 +DOG_TORSO_Z = 0.40 + + +def _quadrotor( + prefix: str, x: float, y: float, z: float = 0.15, accent: str = "0.95 0.26 0.21 1" +) -> tuple[str, str, str]: + """One PX4-compatible quadrotor body.""" + rotors, thrusts, motors, props = [], [], [], [] + have_mesh = x500_meshes_available() + for i, (rx, ry, rgba) in enumerate(ROTORS): + # The disc is the collision proxy for a spinning prop. It stays in the + # model when the meshes are on, just hidden (group 3). + rotors.append( + f' ' + ) + if have_mesh: + props.append( + f' \n' + f' \n' + # Spinning props read as a translucent disc, not as blades. The + # bridge cross-fades prop mesh -> this disc with throttle, the + # same trick Gazebo's iris uses. Alpha 0 here: invisible until + # the motors actually spin. + f' ' + ) + thrusts.append( + f' ' + ) + motors.append( + f' ' + ) + legs = "\n".join( + f' ' + for i, (sx, sy) in enumerate([(1, 1), (1, -1), (-1, 1), (-1, -1)]) + ) + # With the X500 meshes present the primitives stay in the model as the + # COLLISION shapes but move to group 3, which the viewer hides by default + # (press 3 in the viewer to see them). The meshes are pure decoration: + # contype/conaffinity 0 and mass 0, so mass, inertia and every contact are + # byte-identical either way. Toggling the meshes cannot change how it flies. + if x500_meshes_available(): + # The recognisable X500 details the frame scan does not carry: the GPS + # mast with its white puck, the battery slab, and a rear status LED the + # bridge drives (green = disarmed, red = armed). All visual-only. + dressing = f""" + + + + """ + vis_group, mesh_visual = "3", f""" + +{chr(10).join(props)}""" + else: + dressing = "" + vis_group, mesh_visual = "0", "" + + body = f""" + + + {mesh_visual} + {dressing} + +{chr(10).join(rotors)} +{legs} + +{chr(10).join(thrusts)} + """ + sensors = f""" + """ + return body, "\n".join(motors), sensors + + +def _quadruped( + prefix: str, x: float, y: float, accent: str = "1.00 0.76 0.03 1" +) -> tuple[str, str, str]: + """One mesh-free quadruped: box torso, four 2-DOF legs, position servos.""" + legs, actuators = [], [] + corners = [("fl", 0.15, 0.11), ("fr", 0.15, -0.11), ("rl", -0.15, 0.11), ("rr", -0.15, -0.11)] + for name, lx, ly in corners: + legs.append(f""" + + + + + + + + """) + # Position servos hold the stance. `armature` is the important term: + # it models geared-motor rotor inertia, and without it a stiff actuator + # on a light link is only marginally stable at a 4 ms step -- the stance + # height then depends non-monotonically on kp and MuJoCo emits NaN QACC + # warnings. With armature the result is identical across a 10x gain + # range, so the gain can stay low and gentle. + actuators.append( + f' ' + ) + actuators.append( + f' ' + ) + body = f""" + + + + + + + +{chr(10).join(legs)} + """ + return body, "\n".join(actuators), "" + + +def build_scene(n_drones: int = 3, n_dogs: int = 1, timestep: float = 0.004) -> str: + """Return MJCF for a world holding *n_drones* quadrotors and *n_dogs* quadrupeds. + + Drone spawn poses mirror the Gazebo swarm layout (a 3 m triangle) so the + demo geometry, and therefore the separation guardrail's expectations, carry + over from the Gazebo runs unchanged. + """ + drone_poses = [(0.0, 2.0), (-2.0, -1.0), (2.0, -1.0)] + bodies, actuators, sensors = [], [], [] + spawned: list[tuple[float, float]] = [] + + for i in range(n_drones): + if i < len(drone_poses): + px, py = drone_poses[i] + else: + # Drones beyond the demo triangle go on a 3 m grid north of it, + # four per row. The previous "ring" formula (px + 4*(1 + i//3)) + # collided with itself: drones 6 and 8 both landed on (10, -1), + # spawning inside each other. The pairwise assertion below now + # proves whatever this produces, at every fleet size. + k = i - len(drone_poses) + px = -4.5 + 3.0 * (k % 4) + py = 4.5 + 3.0 * (k // 4) + spawned.append((px, py)) + b, a, s = _quadrotor(f"drone{i}", px, py, accent=_accent(i)) + bodies.append(b) + actuators.append(a) + sensors.append(s) + + # Landing pads mark each drone's home. Visual only (contype/conaffinity 0) + # so they can never perturb a touchdown, but they give the viewer a fixed + # ground reference and make an RTL legible at a glance. + pads: list[str] = [] + for i, (px, py) in enumerate(spawned): + pads.append( + f' \n' + f' \n' + f' ' + ) + + dog_poses = [(-4.0 - 1.5 * i, 0.0) for i in range(n_dogs)] + for i, (dx, dy) in enumerate(dog_poses): + b, a, _ = _quadruped(f"dog{i}", dx, dy, accent=_accent(i + 2)) + bodies.append(b) + actuators.append(a) + + # Spawn sanity: robots that spawn inside each other produce contact forces + # on the very first step that read as an inexplicable launch-and-tumble. + # The ring formula for >3 drones is ad hoc, so prove the poses rather than + # trust them -- at every fleet size, not just the ones tried so far. + all_poses = spawned + dog_poses + for i_a in range(len(all_poses)): + for i_b in range(i_a + 1, len(all_poses)): + ax, ay = all_poses[i_a] + bx, by = all_poses[i_b] + d = ((ax - bx) ** 2 + (ay - by) ** 2) ** 0.5 + if d < 1.2: + raise ValueError( + f"spawn poses {i_a} ({ax:.1f},{ay:.1f}) and {i_b} " + f"({bx:.1f},{by:.1f}) are {d:.2f} m apart -- robots would " + "spawn overlapping. Fix the pose table in build_scene()." + ) + + return f""" + + + +""" + + +if __name__ == "__main__": + import argparse + + p = argparse.ArgumentParser(description="Emit a mixed drone/legged MuJoCo scene") + p.add_argument("--drones", type=int, default=3) + p.add_argument("--dogs", type=int, default=1) + p.add_argument("-o", "--out", default="-") + a = p.parse_args() + xml = build_scene(a.drones, a.dogs) + if a.out == "-": + print(xml) + else: + with open(a.out, "w") as fh: + fh.write(xml) + print(f"wrote {a.out}") + + +# --------------------------------------------------------------------------- +# Real Unitree Go1 composition +# --------------------------------------------------------------------------- +# The primitive quadruped above runs anywhere with no assets. When MuJoCo +# Menagerie is present we can do much better: attach the real Go1 that DimOS's +# trained locomotion policy was trained against, one per dog, each under its own +# `dogN-` prefix. Same `MjSpec.attach(prefix=...)` call DimOS's MujocoSimModule +# uses -- just in a loop, because this world holds a whole fleet rather than one +# robot. +# +# Note the Go1 collides on FEET ONLY (4 geoms). That is not a shortcut: it is the +# model the policy was trained on, and it is cheaper than the primitive dog, +# which collides on feet, calves and thighs. + +GO1_XML_REL = "mujoco_playground/_src/locomotion/go1/xmls/go1_mjx_feetonly.xml" +GO1_SPAWN_Z = 0.30 # menagerie "home" keyframe stands at z=0.27 +GO1_HOME_ANGLES = [0.0, 0.9, -1.8] * 4 # from unitree_go1/go1.xml keyframe + + +X500_MESH_DIR = Path(__file__).parent / "assets" / "x500" +X500_MESHES = ("frame", "motor_base", "motor_bell", "prop_cw", "prop_ccw") +X500_MESH_EXT = ".stl" + + +def x500_meshes_available() -> bool: + """True when the decimated X500 meshes are on disk. + + Fetched and converted by ``tools/fetch_x500_meshes.py``. They are optional: + without them the drones fall back to the box-and-cylinder airframe, which + is dimensionally identical and flies exactly the same. + """ + return all((X500_MESH_DIR / f"{m}{X500_MESH_EXT}").exists() for m in X500_MESHES) + + +def _x500_asset_block() -> str: + if not x500_meshes_available(): + return "" + return "\n".join( + f' ' + for m in X500_MESHES + ) + + +def _environment() -> str: + """Static dressing for the demo field. Frames: MuJoCo NWU, so y = -East. + + Three kinds of object, with different physics rules on purpose: + * field lines -- visual only, mark the default 40 x 30 m sweep area + (NED corners (0,0)..(40,30)) so a grid_sweep is legible in the viewer; + * trees -- visual only, parked well outside the field. Purely for depth + perception; a visual-only trunk the dog could walk through must never + sit anywhere a mission path goes; + * crates -- the one COLLIDABLE cluster, at NED (12, -14), i.e. west of + the field and away from spawns and mission routes. They exist so there + is at least one honest obstacle in the world; nothing routes through + them by default. + """ + out: list[str] = [] + # Field lines: NED (0,0)-(40,30) -> MuJoCo x 0..40, y -30..0. + lw, lh = 0.12, 0.012 + for name, x, y, sx, sy in ( + ("field_s", 0.0, -15.0, lw, 15.0), + ("field_n", 40.0, -15.0, lw, 15.0), + ("field_w", 20.0, 0.0, 20.0, lw), + ("field_e", 20.0, -30.0, 20.0, lw), + ): + out.append( + f' ' + ) + trees = [(-25, 18), (-32, -14), (15, 22), (52, 9), (-14, -26), (34, 20), (55, -18)] + for i, (x, y) in enumerate(trees): + out.append( + f' ' + ) + out.append( + f' ' + ) + crates = [(12.0, 14.0, 0.30), (13.1, 14.7, 0.24), (12.4, 15.5, 0.20)] + for i, (x, y, h) in enumerate(crates): + out.append( + f' ' + ) + return "\n".join(out) + + +def go1_unavailable_reason() -> str | None: + """Why the real Go1 cannot be used, or None when it can. + + Falling back to the primitive quadruped silently is worse than failing: the + robot still walks, just badly, and there is nothing on screen or in the log + to say a much better one was available. Every caller reports this string. + """ + try: + from mujoco_playground._src import mjx_env + except Exception as exc: + return f"mujoco_playground not importable ({exc}); pip install playground" + + import sysconfig + + candidate = Path(sysconfig.get_paths()["purelib"]) / GO1_XML_REL + if not candidate.exists(): + return f"Go1 MJCF missing at {candidate}" + # The MJCF references menagerie by a relative path; check the assets are + # really there rather than failing deep inside the MuJoCo compiler. + assets = mjx_env.MENAGERIE_PATH / "unitree_go1" / "assets" + if not assets.exists(): + return ( + f"MuJoCo Menagerie not downloaded ({assets}); run: python -c " + '"from mujoco_playground._src import mjx_env; mjx_env.ensure_menagerie_exists()"' + ) + return None + + +def go1_xml_path() -> str | None: + """Absolute path to the Go1 MJCF the policy was trained on, or None.""" + if go1_unavailable_reason() is not None: + return None + import sysconfig + + return str(Path(sysconfig.get_paths()["purelib"]) / GO1_XML_REL) + + +def build_model(n_drones: int = 3, n_dogs: int = 1, timestep: float = 0.004, + real_go1: bool = True) -> tuple[Any, bool]: + """Compile the fleet world. Returns (MjModel, used_real_go1). + + With Menagerie available every dog is a real Go1 attached under ``dogN-``. + Without it we fall back to the primitive quadruped so the simulator still + runs on a machine that has no assets. + """ + import mujoco + + go1 = go1_xml_path() if (real_go1 and n_dogs > 0) else None + if go1 is None: + return mujoco.MjModel.from_xml_string(build_scene(n_drones, n_dogs, timestep)), False + + # Drones and world only; dogs get attached as real Go1 bodies below. + spec = mujoco.MjSpec.from_string(build_scene(n_drones, 0, timestep)) + for i in range(n_dogs): + child = mujoco.MjSpec.from_file(go1) + # The Go1 MJCF asks for Euler + 1 solver iteration (what the policy was + # trained under). This world keeps implicitfast + 100 iterations, which + # the drones need for stability. Adopt the parent's settings explicitly + # instead of letting attach() resolve the conflict and warn on every + # startup -- the policy was measured to transfer cleanly across the + # difference (1.0 m/s commanded -> 0.94 m/s achieved, no falls). + child.option.integrator = spec.option.integrator + child.option.iterations = spec.option.iterations + child.option.ls_iterations = spec.option.ls_iterations + child.option.timestep = spec.option.timestep + frame = spec.worldbody.add_frame(pos=[-4.0 - 1.5 * i, 0.0, GO1_SPAWN_Z]) + spec.attach(child, prefix=f"dog{i}-", frame=frame) + model = spec.compile() + model.opt.timestep = timestep + return model, True + + +def go1_policy_path() -> str | None: + """Path to DimOS's trained Go1 locomotion policy, or None if unavailable. + + Ships in the Git-LFS archive ``data/.lfs/mujoco_sim.tar.gz``; ``get_data`` + unpacks it on first use. Menagerie supplies the mesh but NOT the weights, so + both have to be present before the real robot can walk. + """ + try: + from dimos.utils.data import get_data + + path = Path(str(get_data("mujoco_sim"))) / "unitree_go1_policy.onnx" + except Exception as exc: + logger.warning( + f"trained Go1 policy unavailable ({exc}). " + "Needs git-lfs: sudo apt install git-lfs && git lfs pull " + '--include "data/.lfs/mujoco_sim.tar.gz"' + ) + return None + if not path.exists(): + logger.warning(f"trained Go1 policy not found at {path}") + return None + return str(path) diff --git a/dimos/simulation/px4_hil/sim.sh b/dimos/simulation/px4_hil/sim.sh new file mode 100755 index 0000000000..e5b555df69 --- /dev/null +++ b/dimos/simulation/px4_hil/sim.sh @@ -0,0 +1,211 @@ +#!/usr/bin/env bash +# DimOS mixed-fleet simulator: MuJoCo physics + PX4 SITL + the DimOS stack. +# +# ./sim.sh start [DRONES] [DOGS] [--viewer] +# bring everything up (default 3 1) +# --viewer opens the MuJoCo window +# ./sim.sh stop tear everything down +# ./sim.sh status what is running +# ./sim.sh log [bridge|daemon|px4N] +# +# Order matters and is handled here: the MuJoCo bridge must be listening before +# PX4 dials in on 4560+i, PX4 must be up before its simulated battery can be +# tuned, and the DimOS daemon must come last so every vehicle is already there. +set -u + +HERE="$(cd "$(dirname "$0")" && pwd)" +REPO="$(cd "$HERE/../../.." && pwd)" +PX4_DIR="${PX4_DIR:-$HOME/PX4-Autopilot}" +PX4_BUILD="$PX4_DIR/build/px4_sitl_default" +VENV="$REPO/.venv/bin" +LOGS="${SIM_LOG_DIR:-/tmp/dimos-sim}" +# Match the bridge by its module basename only. A pattern like +# 'dimos.simulation.px4_hil' looks safe but `.` is a regex wildcard, so it also +# matches this script's own path (dimos/simulation/px4_hil/sim.sh) and the +# teardown kills itself before it can start anything. +PATTERN='fleet_bridge' + +mkdir -p "$LOGS" + +_stop() { + "$VENV/dimos" stop >/dev/null 2>&1 || true + for pid in $(pgrep -f "$PATTERN" 2>/dev/null); do + [ "$pid" = "$$" ] && continue + kill -9 "$pid" 2>/dev/null || true + done + for pid in $(pgrep -x px4 2>/dev/null); do kill -9 "$pid" 2>/dev/null || true; done + rm -f /tmp/px4-sock-* 2>/dev/null || true + sleep 2 + echo "stopped." +} + +_status() { + echo "px4 instances : $(pgrep -xc px4 2>/dev/null || echo 0)" + echo "mujoco bridge : $(pgrep -fc "$PATTERN" 2>/dev/null || echo 0)" + "$VENV/dimos" status 2>&1 | head -4 + echo "governor : $(cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor 2>/dev/null || echo '?')" + last=$(grep -ao 'sim=[ 0-9.]*s rtf=[ 0-9.]*x' "$LOGS/bridge.log" 2>/dev/null | tail -1) + echo "physics : ${last:-no rtf line yet (starting, or bridge down)}" +} + +_start() { + local drones="${1:-3}" dogs="${2:-1}" viewer="${3:-}" + case "$drones$dogs" in + *[!0-9]*) echo "usage: start [DRONES] [DOGS] [--viewer] -- counts must be integers, got '$drones' '$dogs'" >&2; exit 1 ;; + esac + if [ "$drones" -eq 0 ] && [ "$dogs" -eq 0 ]; then + echo "nothing to simulate: pass at least one drone or dog" >&2; exit 1 + fi + if [ "$drones" -gt 5 ] || [ "$dogs" -gt 5 ]; then + echo "refusing $drones drones / $dogs dogs: the operating envelope is 5 drones + 5 dogs." >&2 + echo "Everything past that is unvalidated. Edit sim.sh if the envelope changes." >&2; exit 1 + fi + gov=$(cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor 2>/dev/null || echo '?') + if [ "$gov" = "powersave" ]; then + echo "WARNING: CPU governor is 'powersave' (~1.5 GHz of a 4.7 GHz machine)." >&2 + echo " This costs ~2.4x realtime. Fix: sudo cpupower frequency-set -g performance" >&2 + fi + [ -x "$PX4_BUILD/bin/px4" ] || { echo "PX4 not built at $PX4_BUILD" >&2; exit 1; } + + _stop >/dev/null + if [ "$drones" -gt 0 ]; then + for i in $(seq 0 $((drones - 1))); do + port=$((4560 + i)) + if ss -ltn 2>/dev/null | grep -q ":$port "; then + echo "port $port is still in use after teardown:" >&2 + ss -ltnp 2>/dev/null | grep ":$port " >&2 || true + echo "kill that process, then retry." >&2 + exit 1 + fi + done + fi + rm -f "$LOGS"/*.log + + if [ -n "$viewer" ] && [ -z "${DISPLAY:-}" ] && [ -z "${WAYLAND_DISPLAY:-}" ]; then + echo " (--viewer requested but DISPLAY and WAYLAND_DISPLAY are both unset;" >&2 + echo " no window can open. Run from a desktop session.)" >&2 + fi + echo "[1/4] MuJoCo world: $drones drone(s), $dogs dog(s)${viewer:+ (viewer)}" + ( cd "$REPO" && setsid "$VENV/python" -m dimos.simulation.px4_hil.fleet_bridge \ + --drones "$drones" --dogs "$dogs" $viewer >"$LOGS/bridge.log" 2>&1 /dev/null 2>&1; then + echo "" >&2 + echo "MuJoCo bridge failed to start. Last lines of $LOGS/bridge.log:" >&2 + tail -20 "$LOGS/bridge.log" >&2 + _stop >/dev/null + return 1 + fi + + if [ "$drones" -gt 0 ]; then + echo "[2/4] PX4 SITL x$drones" + for i in $(seq 0 $((drones - 1))); do + # Start from a clean parameter store. A stale or corrupt parameters.bson + # produces bizarre estimator behaviour that looks like a swarm bug. + rm -f "$PX4_BUILD/rootfs/$i"/parameters*.bson 2>/dev/null || true + ( cd "$PX4_BUILD" && PX4_SIM_MODEL=none_iris setsid ./bin/px4 -i "$i" -d \ + >"$LOGS/px4_$i.log" 2>&1 empty in 60 s)" + if ! ( cd "$REPO" && "$VENV/python" \ + dimos/simulation/px4_hil/sim_params.py --count "$drones" ) \ + >"$LOGS/params.log" 2>&1; then + echo " first attempt failed, retrying once..." + if ! ( cd "$REPO" && "$VENV/python" \ + dimos/simulation/px4_hil/sim_params.py --count "$drones" ) \ + >>"$LOGS/params.log" 2>&1; then + echo "sim parameter tuning FAILED twice -- without it batteries drain in 60 s" >&2 + echo "and the geofence is not armed. See $LOGS/params.log" >&2 + _stop >/dev/null # do not leave a half-started stack behind + exit 1 + fi + fi + else + echo "[2/4] no drones requested, skipping PX4" + echo "[3/4] no battery to tune" + fi + + # SIM_BLUEPRINT=mixed-fleet-agentic adds the natural-language agent on top + # of the same fleet (needs OPENAI_API_KEY exported in this shell; drive it + # with `dimos humancli`). Default is the plain MCP blueprint. + blueprint="${SIM_BLUEPRINT:-mixed-fleet-mcp}" + echo "[4/4] DimOS daemon ($blueprint, SIM_DRONES=$drones SIM_DOGS=$dogs)" + if [ "$blueprint" != "mixed-fleet-mcp" ] && [ -z "${OPENAI_API_KEY:-}" ]; then + echo "WARNING: OPENAI_API_KEY is not set -- the agent in $blueprint will not work." >&2 + fi + ( cd "$REPO" && CI=1 DIMOS_TRANSPORT=zenoh SIM_DRONES="$drones" SIM_DOGS="$dogs" \ + setsid "$VENV/dimos" run "$blueprint" >"$LOGS/daemon.log" 2>&1 /dev/null \ + | grep -cE "^ (drone|dog)[0-9]+ .*connected=True" || true) + if ! pgrep -f "$PATTERN" >/dev/null 2>&1; then + printf '\n' + echo "MuJoCo bridge died during startup. Last lines of $LOGS/bridge.log:" >&2 + tail -20 "$LOGS/bridge.log" >&2 + _stop >/dev/null + return 1 + fi + if [ "${got:-0}" -ge "$want" ]; then + # Final gate: the physics clock must be moving. Everything above can pass + # while the world is frozen. + # Sample either side of a full status-log period. The bridge prints + # "sim=..." every 5 s, so a 2 s window reads the SAME line twice and + # declares a perfectly healthy world frozen. + local t0 t1 + t0=$(grep -ao 'sim=[ 0-9.]*' "$LOGS/bridge.log" | tail -1 | tr -dc '0-9.') + sleep 7 + t1=$(grep -ao 'sim=[ 0-9.]*' "$LOGS/bridge.log" | tail -1 | tr -dc '0-9.') + if [ -n "$t0" ] && [ -n "$t1" ] && [ "$t0" = "$t1" ]; then + printf '\n' + echo "robots are connected but the sim clock is not advancing" >&2 + echo "(stuck at sim=${t0}s). Check $LOGS/bridge.log" >&2 + _stop >/dev/null + return 1 + fi + printf ' ok\n\n' + timeout 20 "$VENV/dimos" mcp call list_drones 2>/dev/null + printf '\nready. logs in %s\n' "$LOGS" + return 0 + fi + printf '.' + sleep 3 + done + printf '\n' + echo "timed out waiting for $want robot(s); check $LOGS/daemon.log" >&2 + echo "(stack left RUNNING for diagnosis -- sim.sh stop when done)" >&2 + return 1 +} + +case "${1:-}" in + start) + shift + _viewer=""; _drones=""; _dogs="" + for _a in "$@"; do + if [ "$_a" = "--viewer" ]; then _viewer="--viewer" + elif [ -z "$_drones" ]; then _drones="$_a" + elif [ -z "$_dogs" ]; then _dogs="$_a" + fi + done + _start "${_drones:-3}" "${_dogs:-1}" "$_viewer" ;; + stop) _stop ;; + status) _status ;; + log) tail -f "$LOGS/${2:-bridge}.log" ;; + *) echo "usage: $0 {start [DRONES] [DOGS] [--viewer]|stop|status|log [bridge|daemon|px4N]}" >&2; exit 1 ;; +esac diff --git a/dimos/simulation/px4_hil/sim_params.py b/dimos/simulation/px4_hil/sim_params.py new file mode 100644 index 0000000000..10ee94dcbd --- /dev/null +++ b/dimos/simulation/px4_hil/sim_params.py @@ -0,0 +1,195 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Make PX4 SITL instances survive a full demo run. + +PX4's simulated battery drains from 100% to empty in ``SIM_BAT_DRAIN`` seconds +— and the default is **60**. Any demo longer than about a minute (a grid sweep +easily is) trips the low-battery failsafe mid-flight: the drone abandons its +mission, tries to descend while OFFBOARD setpoints are still streaming, fights +itself, and usually ends up inverted on the ground hundreds of metres away. + +That failure looks exactly like a swarm bug. It is not. This script stretches +the simulated battery so a demo can actually finish. + +These are SIMULATION-ONLY parameters (``SIM_*``). They do not exist on real +hardware, so nothing here can mask a real low-battery condition in the field. + +Run after the instances are up, before pointing DimOS at them: + + .venv/bin/python dimos/simulation/px4_hil/sim_params.py --count 3 +""" + +from __future__ import annotations + +import argparse +import struct +import sys +import time + +# Simulation-only battery behaviour. Held long enough that a full demo +# (takeoff -> sweep -> formation -> RTL) never trips a low-battery failsafe. +PARAMS: dict[str, float] = { + # SIM-seconds, like every PX4 time param (the sim clock runs 10-18x wall). + # 45000 sim-s ~= 40-75 wall-minutes full->empty: outlasts any demo without + # pretending batteries are infinite. + "SIM_BAT_DRAIN": 45000.0, + "SIM_BAT_MIN_PCT": 60.0, # never report below this + # --- X500 airframe geometry ------------------------------------------- + # We boot `none_iris` because it is the only HIL-capable quadrotor airframe + # (every x500 airframe is `gz_`-prefixed and starts gz_bridge, which this + # simulator does not use). So the plumbing is iris and the geometry is + # overridden here to a real X500: a symmetric 0.174 m square, matching the + # x500_base SDF the visual mesh comes from. + # + # These MUST stay in lockstep with ROTORS in dimos/simulation/px4_hil/scene.py. + # PX4 builds its control-allocation matrix from these positions; if MuJoCo + # puts the thrust somewhere else, the controller mis-allocates torque and + # the vehicle flips on arming. + # + # Note PX4's own 4001_gz_x500 uses (0.13, 0.22), which does NOT match its + # own x500 mesh (0.174, 0.174). We use the true airframe geometry so PX4, + # MuJoCo and the mesh all agree. + "CA_ROTOR0_PX": 0.174, + "CA_ROTOR0_PY": 0.174, + "CA_ROTOR1_PX": -0.174, + "CA_ROTOR1_PY": -0.174, + "CA_ROTOR2_PX": 0.174, + "CA_ROTOR2_PY": -0.174, + "CA_ROTOR3_PX": -0.174, + "CA_ROTOR3_PY": 0.174, + # --- geofence: the in-flight backstop -------------------------------- + # The coordinator's dispatch-time operating radius (max_range_m=250) can + # only judge *commands*. A drone drifting under set_velocity, or a + # maneuver interrupted mid-flight, is invisible to it. PX4's own geofence + # is the layer that sees actual position: at 260 m horizontal / 45 m + # vertical it switches to hold. 260 not 250 so the dispatch gate always + # trips first and the operator sees the polite refusal, not the fence. + "GF_ACTION": 2.0, # 2 = hold position at the fence + "GF_MAX_HOR_DIST": 260.0, + "GF_MAX_VER_DIST": 45.0, + # --- datalink-loss failsafe ------------------------------------------ + # The "laptop dies mid-flight" case. With the default (0 = disabled) the + # fleet just keeps hovering forever when DimOS goes away. 2 = return to + # launch and land, on the autopilot's own authority, no ground station + # involved. This is exactly the behaviour wanted on real hardware, so the + # sim should exercise it too. + "NAV_DLL_ACT": 2.0, + # COM_DL_LOSS_T is measured in PX4's OWN clock -- which this simulator + # drives at 10-18x wall speed. Our GCS heartbeats tick at 1 Hz of WALL + # time, i.e. one every 10-18 sim-seconds, so against the 10 s default the + # connection expired between every pair of beats: "GCS connection + # regained" once a second, arming a lottery, and mid-flight failsafes with + # a live link. (PX4's rcS scales this by PX4_SIM_SPEED_FACTOR for exactly + # this reason.) 300 sim-seconds ~= 15-30 wall-seconds at our speeds: the + # laptop dying still brings the fleet home in under half a minute of wall + # time, and a 1 Hz heartbeat is never mistaken for a dead link. + "COM_DL_LOSS_T": 300.0, + # Same sim-clock arithmetic as COM_DL_LOSS_T, for OFFBOARD: the default + # 1.0 (sim-)second staleness window meets our wall-clock setpoint stream + # (20 Hz wall = ~0.4-0.9 sim-s gaps at 10-18x) on a knife edge. 30 sim-s + # keeps the protection while making it immune to realtime-factor swings. + "COM_OF_LOSS_T": 30.0, +} + +# Parameters PX4 declares as INT32. MAVLink transports these BIT-CAST into the +# float field, not converted: sending a REAL32 2.0 makes PX4 store 1073741824 +# (the bits of 2.0f), which for a bounded param like GF_ACTION is out of range +# and silently rejected -- no geofence. Conversely, sending INT32 for a FLOAT +# param (SIM_BAT_DRAIN is FLOAT) is ignored just as silently, so this set must +# list exactly PX4's INT32 params and nothing else. Verified live both ways. +INT32_PARAMS = {"GF_ACTION", "NAV_DLL_ACT", "COM_DL_LOSS_T"} + +DEFAULT_PORT_BASE = 14540 +CONNECT_TIMEOUT_S = 20.0 +VERIFY_TIMEOUT_S = 5.0 + + +def tune_instance(instance: int, host: str, port_base: int) -> bool: + """Set the sim params on one PX4 instance. Returns True if all verified.""" + from pymavlink import mavutil # type: ignore[import-untyped] + + endpoint = f"udp:{host}:{port_base + instance}" + print(f"[sim-params] instance {instance}: connecting to {endpoint}") + conn = mavutil.mavlink_connection(endpoint, source_system=245, source_component=190) + + if conn.wait_heartbeat(timeout=CONNECT_TIMEOUT_S) is None: + print(f"[sim-params] instance {instance}: NO HEARTBEAT — skipped", file=sys.stderr) + return False + + ok = True + for name, value in PARAMS.items(): + if name in INT32_PARAMS: + # Pack the integer's bytes into the float field (see INT32_PARAMS). + wire_value = struct.unpack(" 1e-3: + print( + f"[sim-params] instance {instance}: {name} readback {confirmed} != {value}", + file=sys.stderr, + ) + ok = False + else: + print(f"[sim-params] instance {instance}: {name} = {confirmed:g}") + + conn.close() + return ok + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--count", type=int, default=3, help="number of SITL instances") + ap.add_argument("--host", default="127.0.0.1") + ap.add_argument("--port-base", type=int, default=DEFAULT_PORT_BASE) + args = ap.parse_args() + + failures = [i for i in range(args.count) if not tune_instance(i, args.host, args.port_base)] + if failures: + print(f"[sim-params] FAILED on instance(s): {failures}", file=sys.stderr) + return 1 + print(f"[sim-params] all {args.count} instance(s) tuned.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/fetch_x500_meshes.py b/tools/fetch_x500_meshes.py new file mode 100755 index 0000000000..f22261c60f --- /dev/null +++ b/tools/fetch_x500_meshes.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fetch the Holybro X500 meshes and convert them for MuJoCo. + + .venv/bin/python tools/fetch_x500_meshes.py + +Source is PX4/PX4-gazebo-models (BSD-3-Clause), the same assets Gazebo renders, +so the simulated drone looks like the one people are used to seeing there. + +Two conversions are needed and neither is optional: + +* **MuJoCo cannot load COLLADA.** It reads STL, OBJ and MSH. The X500 frame and + motors ship as ``.dae``, so they are converted to OBJ here (needs pycollada). +* **The frame must NOT be decimated.** It is ~123,000 disconnected pieces of + triangle soup (standoffs, screws, plates), not a welded surface. Decimating it + to 16k destroyed 98% of its surface area and left a scatter of fragments that + rendered as a few specks in mid-air. Welding first does not save it either: + 60k still loses 55% of the area. MuJoCo stores each mesh ONCE and instances it + per drone, so the full 179k costs the same whether you fly one drone or six. + Binary STL keeps it to ~9 MB. The props are solid single parts and do decimate + cleanly (89% of area at 4k faces). + +Every mesh is also **centred on its bounding box**. Left uncentred, MuJoCo +re-centres a mesh geom and shifts ``geom_pos`` by the centroid, which threw every +prop (0.026, 0.173) off the arm it was supposed to sit on. + +The meshes are purely decorative: ``scene.py`` attaches them with contype 0 and +mass 0, keeping the primitives as the collision shapes. If this script is never +run, the simulator falls back to the box-and-cylinder airframe and flies exactly +the same. +""" + +from __future__ import annotations + +from pathlib import Path +import sys +import urllib.request + +BASE = "https://raw.githubusercontent.com/PX4/PX4-gazebo-models/main/models/x500_base/meshes" +OUT = Path(__file__).resolve().parents[1] / "dimos/simulation/px4_hil/assets/x500" + +# source file -> (output name, target triangle count) +# source file -> (output name, target triangle count; None = do not decimate) +MESHES: dict[str, tuple[str, int | None]] = { + "NXP-HGD-CF.dae": ("frame.stl", None), + "5010Base.dae": ("motor_base.stl", None), + "5010Bell.dae": ("motor_bell.stl", None), + "1345_prop_ccw.stl": ("prop_ccw.stl", 4000), + "1345_prop_cw.stl": ("prop_cw.stl", 4000), +} + + +def main() -> int: + try: + import trimesh + except ImportError: + print("needs trimesh: uv pip install trimesh pycollada fast_simplification", file=sys.stderr) + return 1 + + OUT.mkdir(parents=True, exist_ok=True) + tmp = OUT / ".download" + tmp.mkdir(exist_ok=True) + + total = 0.0 + for src, (dst, target) in MESHES.items(): + raw = tmp / src + if not raw.exists(): + print(f" fetching {src} ...", flush=True) + try: + urllib.request.urlopen(f"{BASE}/{src}", timeout=120) + except Exception as exc: + print(f" FAILED: {exc}", file=sys.stderr) + return 1 + urllib.request.urlretrieve(f"{BASE}/{src}", raw) + + mesh = trimesh.load(raw, force="mesh") + before, area_before = len(mesh.faces), mesh.area + if target and before > target: + mesh.merge_vertices() + mesh = mesh.simplify_quadric_decimation(face_count=target) + mesh.apply_translation(-mesh.bounding_box.centroid) + # Strip materials and UVs; MuJoCo colours these from the MJCF. + mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces) + mesh.export(OUT / dst) + mb = (OUT / dst).stat().st_size / 1e6 + total += mb + kept = mesh.area / area_before * 100 if area_before else 100.0 + print( + f" {src:22s} {before:>8,} -> {len(mesh.faces):>7,} tris " + f"area {kept:5.1f}% kept {dst:15s} {mb:5.2f} MB" + ) + if kept < 80.0: + print(" WARNING: lost most of the surface -- do not decimate this one", file=sys.stderr) + + for leftover in tmp.iterdir(): + leftover.unlink() + tmp.rmdir() + print(f"\nwrote {len(MESHES)} meshes to {OUT} ({total:.2f} MB total)") + print("Source: PX4/PX4-gazebo-models, BSD-3-Clause.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From a49cb4e2bedaec4f354b53c303a483d82fde77cb Mon Sep 17 00:00:00 2001 From: Ezenbaatar Batjargal Date: Mon, 31 Aug 2026 11:38:36 -0700 Subject: [PATCH 2/4] feat(drone): PX4 SITL swarm stack with delivery-verified fleet commands Per-vehicle Px4DroneModule (namespaced), SwarmCoordinator with layered guardrails (operating radius, separation admission, airborne gate, altitude cap), fleet maneuvers (grid sweep, line formation, investigate, staged air-then-ground sweep), and a fleet-level GCS presence that satisfies PX4's arming checks once NAV_DLL_ACT is set. Safety broadcasts repeat 3x and are verified against reported flight modes with per-robot re-send: a lost rtl_all once left an armed drone hovering while its twin landed, and that incident is now a regression test. Co-Authored-By: Claude Fable 5 --- .../basic/drone_px4_sitl_fleet_mcp.py | 91 ++ dimos/robot/drone/mavlink_connection.py | 158 +- dimos/robot/drone/px4_drone_module.py | 557 +++++++ dimos/robot/drone/px4_geo.py | 87 ++ dimos/robot/drone/px4_sitl_connection.py | 732 +++++++++ dimos/robot/drone/px4_sitl_fleet_config.py | 119 ++ dimos/robot/drone/px4_swarm_coordinator.py | 1335 +++++++++++++++++ dimos/robot/drone/test_px4_swarm.py | 1145 ++++++++++++++ dimos/robot/drone/test_px4_swarm_bus.py | 189 +++ 9 files changed, 4391 insertions(+), 22 deletions(-) create mode 100644 dimos/robot/drone/blueprints/basic/drone_px4_sitl_fleet_mcp.py create mode 100644 dimos/robot/drone/px4_drone_module.py create mode 100644 dimos/robot/drone/px4_geo.py create mode 100644 dimos/robot/drone/px4_sitl_connection.py create mode 100644 dimos/robot/drone/px4_sitl_fleet_config.py create mode 100644 dimos/robot/drone/px4_swarm_coordinator.py create mode 100644 dimos/robot/drone/test_px4_swarm.py create mode 100644 dimos/robot/drone/test_px4_swarm_bus.py diff --git a/dimos/robot/drone/blueprints/basic/drone_px4_sitl_fleet_mcp.py b/dimos/robot/drone/blueprints/basic/drone_px4_sitl_fleet_mcp.py new file mode 100644 index 0000000000..614a78dcbb --- /dev/null +++ b/dimos/robot/drone/blueprints/basic/drone_px4_sitl_fleet_mcp.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Namespaced 3-drone PX4 SITL swarm with MCP (no LLM). + +Each vehicle runs its own ``Px4DroneModule`` inside its own namespace, so the +per-drone skills, topics, TF frames and config keys are all separated:: + + drone1/px4dronemodule/takeoff RPC + /drone1/odom topic + -o drone1/px4dronemodule.connection_string=... + +One shared ``SwarmCoordinator`` sits outside the namespaces and owns the +fleet-level surface (state, guardrails, formations, sweeps). Telemetry reaches +it on the exposed ``drone_state`` stream; its commands go back out on the +exposed ``swarm_cmd`` stream. + +Launch the sim first, then drive the whole demo from the CLI — no OpenAI key +needed: + + ./dimos/simulation/px4_hil/sim.sh start 3 0 # in-repo MuJoCo+PX4 simulator + CI=1 dimos run drone-px4-sitl-fleet-mcp --daemon + + dimos mcp modules # three drones + the coordinator + dimos mcp call fleet_state + dimos mcp call takeoff_all --arg altitude=5 + dimos mcp call grid_sweep --arg corner_b_north=40 --arg corner_b_east=30 + dimos mcp call count_within --arg drone=drone1 --arg radius_m=100 + dimos mcp call line_formation + dimos mcp call rtl_all + +Per-drone commands address the instance directly: + + dimos mcp call drone2/px4dronemodule/takeoff --arg altitude=4 +""" + +from dimos.agents.mcp.mcp_server import McpServer +from dimos.core.coordination.blueprints import autoconnect +from dimos.robot.drone.px4_drone_module import Px4DroneModule +from dimos.robot.drone.px4_sitl_fleet_config import get_px4_sitl_fleet_configs +from dimos.robot.drone.px4_swarm_coordinator import SwarmCoordinator + +# Streams that must stay global so they cross the namespace boundary: telemetry +# out of each drone, and fleet commands back in. +FLEET_BUS = {"drone_state", "swarm_cmd"} + +_CONFIGS = get_px4_sitl_fleet_configs() + + +def px4_sitl_swarm( + configs=_CONFIGS, + min_separation_m: float = 2.0, + max_altitude_m: float | None = None, +): + """Compose N namespaced PX4 drones plus one shared coordinator.""" + return autoconnect( + SwarmCoordinator.blueprint( + min_separation_m=min_separation_m, + max_altitude_m=max_altitude_m, + expected_drones=",".join(f"drone{i + 1}" for i in range(len(configs))), + ), + *[ + Px4DroneModule.blueprint( + connection_string=c.connection_string, + instance=c.instance, + sys_id=c.sys_id, + max_altitude_m=max_altitude_m, + ).namespace(f"drone{i + 1}", expose=FLEET_BUS) + for i, c in enumerate(configs) + ], + ) + + +drone_px4_sitl_fleet_mcp = autoconnect( + px4_sitl_swarm(), + McpServer.blueprint(), +).global_config(n_workers=max(2, 2 * len(_CONFIGS))) + +__all__ = ["FLEET_BUS", "drone_px4_sitl_fleet_mcp", "px4_sitl_swarm"] diff --git a/dimos/robot/drone/mavlink_connection.py b/dimos/robot/drone/mavlink_connection.py index e908be4ab0..298a29ee03 100644 --- a/dimos/robot/drone/mavlink_connection.py +++ b/dimos/robot/drone/mavlink_connection.py @@ -17,6 +17,7 @@ import functools import logging +import threading import time from typing import Any @@ -40,6 +41,7 @@ def __init__( connection_string: str = "udp:0.0.0.0:14550", outdoor: bool = False, max_velocity: float = 5.0, + target_system: int | None = None, ) -> None: """Initialize drone connection. @@ -47,10 +49,18 @@ def __init__( connection_string: MAVLink connection string outdoor: Use GPS only mode (no velocity integration) max_velocity: Maximum velocity in m/s + target_system: Optional MAVLink system id to expect on this link. When + provided (as in PX4 SITL multi-vehicle, where each instance maps + to ``sys_id = instance + 1``), ``connect()`` will use this value + instead of waiting for a HEARTBEAT to discover it. This is the + canonical MAVSDK-style multi-vehicle pattern and routes around + PX4 SITL's quirk where Onboard-mode HEARTBEAT only streams to + already-active peers. """ self.connection_string = connection_string self.outdoor = outdoor self.max_velocity = max_velocity + self._expected_target_system = target_system self.mavlink: Any = None # MAVLink connection object self.connected = False self.telemetry: dict[str, Any] = {} @@ -68,21 +78,118 @@ def __init__( # Flag to prevent concurrent fly_to commands self.flying_to_target = False + # GCS heartbeat thread — PX4 SITL refuses to arm without a registered GCS link. + # We satisfy that by impersonating a GCS at 1 Hz (standard MAVLink rate). + self._gcs_hb_stop = threading.Event() + self._gcs_hb_thread: threading.Thread | None = None + def connect(self) -> bool: """Connect to drone via MAVLink.""" try: logger.info(f"Connecting to {self.connection_string}") - self.mavlink = mavutil.mavlink_connection(self.connection_string) - self.mavlink.wait_heartbeat(timeout=30) + # source_system=255 / source_component=190 are the standard "GCS" identifiers, + # so PX4 sees our heartbeats as coming from a ground station and clears its + # "no GCS connection" arming check. + self.mavlink = mavutil.mavlink_connection( + self.connection_string, + source_system=255, + source_component=190, + ) + # Two paths for picking the link's target_system: + # + # (a) `target_system` was passed in (PX4 SITL fleet, RTK swarm, anything + # where the sys_id is known up front). Skip HEARTBEAT discovery — + # wait for *any* MAVLink message from that sys to confirm the link + # is live, then latch the id. This sidesteps PX4 SITL's bug where + # instances >= 1 don't stream HEARTBEAT on their Onboard channel + # until a peer is fully established (chicken-and-egg with our + # own GCS heartbeats). + # + # (b) No explicit target_system. Fall back to HEARTBEAT discovery, + # but skip sys=0 frames (PX4 emits a few during early init before + # MAV_SYS_ID is loaded; latching sys=0 means every command goes + # to broadcast and silently lands nowhere). + deadline = time.time() + 30.0 + picked = False + if self._expected_target_system is not None: + want = self._expected_target_system + while time.time() < deadline: + msg = self.mavlink.recv_match(blocking=True, timeout=1.0) + if msg is None: + continue + if msg.get_srcSystem() == want: + self.mavlink.target_system = want + self.mavlink.target_component = msg.get_srcComponent() or 1 + picked = True + break + if not picked: + logger.error( + f"No MAVLink message from sys={want} within 30s on {self.connection_string}" + ) + return False + else: + while time.time() < deadline: + hb = self.mavlink.recv_match(type="HEARTBEAT", blocking=True, timeout=1.0) + if hb is None: + continue + src = hb.get_srcSystem() + if src and src > 0: + self.mavlink.target_system = src + self.mavlink.target_component = hb.get_srcComponent() or 1 + picked = True + break + if not picked: + logger.error( + f"No HEARTBEAT with sys_id > 0 within 30s on {self.connection_string}" + ) + return False + self.connected = True logger.info(f"Connected to system {self.mavlink.target_system}") + self._start_gcs_heartbeat() self.update_telemetry() return True except Exception as e: logger.error(f"Connection failed: {e}") return False + def _start_gcs_heartbeat(self) -> None: + """Publish a 1 Hz GCS heartbeat on this link, forever. + + The loop is deliberately unkillable: the shared MAVLink connection + races the offboard streamer thread, and an early version that returned + on the first send exception silently disabled arming for the rest of + the session once a datalink-loss failsafe (NAV_DLL_ACT) was + configured. A missed beat at 1 Hz is harmless; a dead loop is not. + + Fleet-level GCS presence on PX4's GCS port (the thing the arming check + actually credits) lives in SwarmCoordinator._gcs_presence_loop; this + per-link heartbeat only marks the offboard channel as alive. + """ + if self._gcs_hb_thread and self._gcs_hb_thread.is_alive(): + return + self._gcs_hb_stop.clear() + + def _loop() -> None: + while not self._gcs_hb_stop.is_set(): + try: + self.mavlink.mav.heartbeat_send( + mavutil.mavlink.MAV_TYPE_GCS, + mavutil.mavlink.MAV_AUTOPILOT_INVALID, + 0, + 0, + 0, + ) + except Exception: + pass # shared-link race; skip this beat, never die + self._gcs_hb_stop.wait(1.0) + + self._gcs_hb_thread = threading.Thread( + target=_loop, name="dimos-mavlink-gcs-hb", daemon=True + ) + self._gcs_hb_thread.start() + def update_telemetry(self, timeout: float = 0.1) -> None: """Update telemetry data from available messages.""" if not self.connected: @@ -547,19 +654,31 @@ def rotate_to(self, target_heading_deg: float, timeout: float = 60.0) -> bool: return False def arm(self) -> bool: - """Arm the drone.""" + """Arm the drone. + + Sends MAV_CMD_COMPONENT_ARM_DISARM and waits for the HEARTBEAT armed + bit to flip. Polls ``self.telemetry["HEARTBEAT"]`` (kept fresh by a + separate telemetry loop, e.g. ``Px4DroneModule._telemetry_loop``) + instead of calling ``recv_match`` directly. ``recv_match`` would race + with the fleet's telemetry consumer for the same MAVLink message + stream — both reading the same socket, each stealing messages from the + other — and that race can silently miss the COMMAND_ACK and HEARTBEAT + frames that confirm arming, returning False even when arming succeeded. + + Returns True only once HEARTBEAT shows the armed bit set; otherwise + False after ``timeout_s``. + """ if not self.connected: return False logger.info("Arming motors...") - self.update_telemetry() self.mavlink.mav.command_long_send( self.mavlink.target_system, self.mavlink.target_component, mavutil.mavlink.MAV_CMD_COMPONENT_ARM_DISARM, 0, - 1, + 1, # param1: 1 = arm 0, 0, 0, @@ -568,24 +687,18 @@ def arm(self) -> bool: 0, ) - # Wait for ACK - ack = self.mavlink.recv_match(type="COMMAND_ACK", blocking=True, timeout=5) - if ack and ack.command == mavutil.mavlink.MAV_CMD_COMPONENT_ARM_DISARM: - if ack.result == mavutil.mavlink.MAV_RESULT_ACCEPTED: - logger.info("Arm command accepted") - - # Verify armed status - for _i in range(10): - msg = self.mavlink.recv_match(type="HEARTBEAT", blocking=True, timeout=1) - if msg: - armed = msg.base_mode & mavutil.mavlink.MAV_MODE_FLAG_SAFETY_ARMED - if armed: - logger.info("Motors ARMED successfully!") - return True - time.sleep(0.5) - else: - logger.error(f"Arm failed with result: {ack.result}") + # Poll the telemetry dict (updated by the consumer thread) for the + # armed flag. 7 s is comfortably above the 1 Hz heartbeat cadence even + # if a few frames are dropped. + deadline = time.time() + 7.0 + while time.time() < deadline: + hb = self.telemetry.get("HEARTBEAT", {}) + if hb.get("armed"): + logger.info("Motors ARMED") + return True + time.sleep(0.1) + logger.warning("Arm timeout — HEARTBEAT armed bit never set within 7s") return False def disarm(self) -> bool: @@ -1004,6 +1117,7 @@ def get_telemetry(self) -> dict[str, Any]: def disconnect(self) -> None: """Disconnect from drone.""" + self._gcs_hb_stop.set() if self.mavlink: self.mavlink.close() self.connected = False diff --git a/dimos/robot/drone/px4_drone_module.py b/dimos/robot/drone/px4_drone_module.py new file mode 100644 index 0000000000..3e44cdbb7e --- /dev/null +++ b/dimos/robot/drone/px4_drone_module.py @@ -0,0 +1,557 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""One PX4 vehicle, one module instance. + +This is the namespaced replacement for the per-drone half of +``Px4SitlFleetModule``. Instead of one module multiplexing N connections behind +``drone="drone-1"`` string arguments, each vehicle gets its own instance under +its own namespace:: + + autoconnect( + *[ + Px4DroneModule.blueprint(connection_string=c.connection_string) + .namespace(f"drone{i + 1}", expose={"drone_state", "swarm_cmd"}) + for i, c in enumerate(get_px4_sitl_fleet_configs()) + ], + ) + +which gives each drone its own RPC surface (``drone1/px4dronemodule/takeoff``), +its own topics (``/drone1/odom``), its own TF frames, and its own config keys +(``-o drone1/px4dronemodule.connection_string=...``). + +Two streams are deliberately *exposed* (left unprefixed, so they stay global and +cross the namespace boundary): + +``drone_state`` + Every drone publishes its telemetry snapshot here. ``SwarmCoordinator`` + subscribes once and sees the whole fleet. +``swarm_cmd`` + The coordinator broadcasts fleet commands here. Each drone acts only on + messages addressed to its own key (or to ``all``). + +Everything else — notably ``cmd_vel`` from a vision tracker — stays namespace +local, so drone 2's tracker can never drive drone 1. +""" + +from __future__ import annotations + +import json +import math +import threading +import time +from typing import Any + +from dimos_lcm.std_msgs import String + +from dimos.agents.annotation import skill +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.robot.drone.px4_sitl_connection import Px4SitlConnection +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +# How often each drone republishes its state snapshot on the shared bus. +STATE_PUBLISH_HZ = 4.0 + + +class Px4DroneConfig(ModuleConfig): + """Per-drone identity and safety envelope. + + Every field is addressable per instance, e.g. + ``-o drone2/px4dronemodule.connection_string=udp:127.0.0.1:14541`` or + ``DRONE2_PX4DRONEMODULE__MAX_ALTITUDE_M=30``. + """ + + # pymavlink connection string. SITL instance N is udp:127.0.0.1:1454N; + # hardware is serial:/dev/ttyACM0:57600 or udp:192.168.1.50:14550. + connection_string: str = "udp:127.0.0.1:14540" + # PX4 SITL instance index (selects the OFFBOARD port PX4 binds locally). + instance: int = 0 + # MAV_SYS_ID this vehicle heartbeats with. Informational for reporting. + sys_id: int = 1 + # Optional ceiling in meters above home (positive up). None = no software + # cap; PX4's own geofence params take over. Set on hardware blueprints. + max_altitude_m: float | None = None + + +class Px4DroneModule(Module): + """A single PX4 vehicle: MAVLink connection, telemetry, and flight skills. + + The skills carry no ``drone`` argument — the instance *is* the drone. When + three of these run under ``drone1``/``drone2``/``drone3`` namespaces the + agent sees three separate tool sets and addresses them by name natively. + """ + + # One vehicle per worker process. Namespacing isolates *names*, not faults: + # without this, modules share workers, so a wedged MAVLink socket or a + # GIL-holding loop on one drone can stall another. Cheap insurance for a + # handful of vehicles, and the same thing GO2Connection does upstream. + dedicated_worker = True + + config: Px4DroneConfig + + # Namespace-local: body-frame velocity from this drone's vision tracker. + # Gated — dropped entirely until start_follow() arms forwarding. + cmd_vel: In[Twist] + + # Namespace-local: this drone's pose, for viewers and TF. + odom: Out[PoseStamped] + + # Exposed (global): telemetry out to the coordinator, commands back in. + drone_state: Out[String] + swarm_cmd: In[String] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.connection: Px4SitlConnection | None = None + self._running = False + self._telemetry_thread: threading.Thread | None = None + self._state_thread: threading.Thread | None = None + # Vision-follow gating: cmd_vel is dropped unless follow is explicitly on. + self._follow_enabled = False + self._follow_lock_altitude = True + + # ------------------------------------------------------------------ + # Identity + # ------------------------------------------------------------------ + + @property + def drone_key(self) -> str: + """Short fleet-facing name for this drone, derived from its namespace. + + Under ``.namespace("drone2")`` the coordinator sets ``instance_name`` to + ``drone2/px4dronemodule``, so the key is ``drone2``. Un-namespaced (a + single-drone blueprint) it falls back to ``drone1``. + """ + name = self.config.instance_name + if name and "/" in name: + return name.rsplit("/", 1)[0] + return name or "drone1" + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + @rpc + def start(self) -> None: + super().start() + self._running = True + + conn = Px4SitlConnection( + connection_string=self.config.connection_string, + instance=self.config.instance, + ) + if not conn.connect(): + logger.warning( + f"[{self.drone_key}] failed to connect at {self.config.connection_string}" + ) + return + self.connection = conn + logger.info( + f"[{self.drone_key}] connected " + f"(sys_id={self.config.sys_id}, endpoint={self.config.connection_string})" + ) + + self._telemetry_thread = threading.Thread( + target=self._telemetry_loop, + daemon=True, + name=f"px4-tel-{self.drone_key}", + ) + self._telemetry_thread.start() + + self._state_thread = threading.Thread( + target=self._state_publish_loop, + daemon=True, + name=f"px4-state-{self.drone_key}", + ) + self._state_thread.start() + + if getattr(self.cmd_vel, "transport", None): + self.cmd_vel.subscribe(self._on_cmd_vel) + logger.info(f"[{self.drone_key}] subscribed to cmd_vel (vision follow input)") + if getattr(self.swarm_cmd, "transport", None): + self.swarm_cmd.subscribe(self._on_swarm_cmd) + logger.info(f"[{self.drone_key}] subscribed to swarm_cmd") + + def stop(self) -> None: + self._running = False + if self.connection is not None: + try: + self.connection.disconnect() + except Exception as e: + logger.debug(f"[{self.drone_key}] disconnect error: {e}") + for t in (self._telemetry_thread, self._state_thread): + if t is not None and t.is_alive(): + t.join(timeout=1.0) + super().stop() + + def _telemetry_loop(self) -> None: + while self._running and self.connection is not None: + try: + self.connection.update_telemetry(timeout=0.05) + except Exception as e: + logger.debug(f"[{self.drone_key}] telemetry error: {e}") + time.sleep(0.1) + + def _state_publish_loop(self) -> None: + """Republish this drone's snapshot on the shared bus at a fixed rate.""" + period = 1.0 / STATE_PUBLISH_HZ + while self._running: + try: + self.drone_state.publish(String(json.dumps(self._state_dict()))) + pose = self._pose_stamped() + if pose is not None: + self.odom.publish(pose) + except Exception as e: + logger.debug(f"[{self.drone_key}] state publish error: {e}") + time.sleep(period) + + # ------------------------------------------------------------------ + # State + # ------------------------------------------------------------------ + + def _state_dict(self) -> dict[str, Any]: + conn = self.connection + if conn is None: + return {"key": self.drone_key, "connected": False, "ts": time.time()} + ned = conn.get_local_ned() + g = conn.get_global_position() + return { + "key": self.drone_key, + # Declared so SwarmCoordinator can pick robots by capability rather + # than assume every reporting robot can fly. Without this a ground + # robot gets handed an air lane in grid_sweep. + "robot_class": "multirotor", + "capabilities": ["air", "camera"], + "connected": True, + "sys_id": self.config.sys_id, + "instance": self.config.instance, + "armed": conn.get_armed(), + "mode": conn.get_mode(), + "battery_pct": conn.get_battery_pct(), + "ned": list(ned) if ned is not None else None, + "altitude_m": (-ned[2]) if ned is not None else None, + "global": ({"lat": g[0], "lon": g[1], "rel_alt_m": g[2]} if g is not None else None), + "ts": time.time(), + } + + def _pose_stamped(self) -> PoseStamped | None: + """This drone's local NED position as a ROS-convention pose (X fwd, Y left, Z up).""" + conn = self.connection + if conn is None: + return None + ned = conn.get_local_ned() + if ned is None: + return None + n, e, d = ned + # MAVLink NED (X=north, Y=east, Z=down) -> ROS/DimOS (X=fwd, Y=left, Z=up). + pose = PoseStamped(ts=time.time(), frame_id=self.frame_id) + pose.position.x = n + pose.position.y = -e + pose.position.z = -d + return pose + + @skill + def state(self) -> str: + """Report this drone's position, armed state, battery, and altitude.""" + return json.dumps(self._state_dict(), indent=2) + + # ------------------------------------------------------------------ + # Guardrails + # ------------------------------------------------------------------ + + def _check_altitude_cap(self, altitude: float) -> str | None: + """Return an error string if ``altitude`` exceeds the configured cap.""" + cap = self.config.max_altitude_m + if cap is None: + return None + if altitude > cap: + return ( + f"{self.drone_key}: REJECTED altitude {altitude:.1f}m exceeds cap " + f"{cap:.1f}m. Stay below the cap or raise max_altitude_m." + ) + return None + + def _require_conn(self) -> Px4SitlConnection | None: + return self.connection + + # ------------------------------------------------------------------ + # Flight skills — no `drone` argument; the instance is the drone + # ------------------------------------------------------------------ + + @skill + def arm(self) -> str: + """Arm this drone's motors.""" + conn = self._require_conn() + if conn is None: + return f"{self.drone_key}: NOT CONNECTED" + return f"{self.drone_key}: arm {'OK' if conn.arm() else 'FAILED'}" + + @skill + def takeoff(self, altitude: float = 3.0) -> str: + """Arm and take off to ``altitude`` (meters above home, positive up).""" + cap_err = self._check_altitude_cap(altitude) + if cap_err: + return cap_err + conn = self._require_conn() + if conn is None: + return f"{self.drone_key}: NOT CONNECTED" + ok = conn.takeoff(altitude) + return f"{self.drone_key}: takeoff to {altitude}m {'OK' if ok else 'FAILED'}" + + @skill + def land(self) -> str: + """Land this drone where it currently is.""" + conn = self._require_conn() + if conn is None: + return f"{self.drone_key}: NOT CONNECTED" + return f"{self.drone_key}: land {'OK' if conn.land() else 'FAILED'}" + + @skill + def rtl(self) -> str: + """Return to launch and land.""" + conn = self._require_conn() + if conn is None: + return f"{self.drone_key}: NOT CONNECTED" + return f"{self.drone_key}: RTL {'OK' if conn.rtl() else 'FAILED'}" + + @skill + def hold(self) -> str: + """Engage AUTO.LOITER (hover in place).""" + conn = self._require_conn() + if conn is None: + return f"{self.drone_key}: NOT CONNECTED" + return f"{self.drone_key}: hold {'OK' if conn.hold() else 'FAILED'}" + + @skill + def goto( + self, + north: float = 0.0, + east: float = 0.0, + altitude: float = 3.0, + yaw_deg: float = 0.0, + ) -> str: + """Fly to a local-NED waypoint via OFFBOARD setpoints. + + Args: + north: North position (m) in this drone's home-relative NED frame. + east: East position (m). + altitude: Altitude above home in meters (positive up). + yaw_deg: Heading in degrees (NED, 0 = north). + + Fleet-wide spacing is enforced by ``SwarmCoordinator``; this skill only + applies the per-drone altitude cap. Prefer the coordinator's + ``investigate``/``grid_sweep`` for multi-drone moves. + """ + cap_err = self._check_altitude_cap(altitude) + if cap_err: + return cap_err + return self._goto_ned(north, east, -altitude, math.radians(yaw_deg)) + + def _goto_ned(self, north: float, east: float, down: float, yaw_rad: float = 0.0) -> str: + conn = self._require_conn() + if conn is None: + return f"{self.drone_key}: NOT CONNECTED" + if not conn._offboard_running and not conn.start_offboard(): + return f"{self.drone_key}: failed to enter OFFBOARD" + conn.set_position_ned(north, east, down, yaw_rad) + return ( + f"{self.drone_key}: goto NED ({north:.1f}, {east:.1f}, " + f"alt={-down:.1f}m), yaw={math.degrees(yaw_rad):.0f}°" + ) + + @skill + def set_velocity( + self, + vn: float = 0.0, + ve: float = 0.0, + vd: float = 0.0, + yaw_rate_deg: float = 0.0, + ) -> str: + """Stream a velocity setpoint (NED, m/s) in OFFBOARD. + + ``vd`` is positive-down. To climb at 1 m/s, pass ``vd=-1.0``. + """ + conn = self._require_conn() + if conn is None: + return f"{self.drone_key}: NOT CONNECTED" + yaw_rate = math.radians(yaw_rate_deg) + if not conn._offboard_running and not conn.start_offboard((vn, ve, vd, yaw_rate)): + return f"{self.drone_key}: failed to enter OFFBOARD" + conn.set_velocity_ned(vn, ve, vd, yaw_rate) + return f"{self.drone_key}: velocity NED ({vn:.2f}, {ve:.2f}, {vd:.2f}) m/s" + + def _follow_path(self, path: list[list[float]], arrival_radius_m: float = 2.0) -> str: + """Execute a waypoint list; the OFFBOARD streamer advances on arrival.""" + conn = self._require_conn() + if conn is None: + return f"{self.drone_key}: NOT CONNECTED" + if not conn._offboard_running and not conn.start_offboard(): + return f"{self.drone_key}: failed to enter OFFBOARD" + conn.set_position_path( + [(p[0], p[1], p[2], p[3] if len(p) > 3 else 0.0) for p in path], + arrival_radius_m=arrival_radius_m, + ) + return f"{self.drone_key}: {len(path)} waypoints dispatched" + + # ------------------------------------------------------------------ + # Emergency + # ------------------------------------------------------------------ + + @skill + def kill(self) -> str: + """Immediately force-disarm this drone — emergency stop. + + Sends MAV_CMD_COMPONENT_ARM_DISARM with the force magic value (21196), + which PX4 honors **even mid-flight**. Motors cut. The drone falls. + **Use ONLY when crashing is preferable to whatever it is about to do** — + e.g. heading toward a person, lost link, autopilot runaway. For normal + shutdown use ``land`` or ``rtl``. + + Also tears down the OFFBOARD streamer so the disarm sticks. + """ + from pymavlink import mavutil + + conn = self._require_conn() + if conn is None: + return f"{self.drone_key}: NOT CONNECTED" + conn.stop_offboard() + conn.mavlink.mav.command_long_send( + conn.mavlink.target_system, + conn.mavlink.target_component, + mavutil.mavlink.MAV_CMD_COMPONENT_ARM_DISARM, + 0, + 0, # param1: 0 = disarm + 21196.0, # param2: force flag (PX4 specific) + 0, 0, 0, 0, 0, + ) + logger.warning(f"[{self.drone_key}] KILL — force-disarm sent") + return f"{self.drone_key}: KILL sent (force-disarm). Motors stopped." + + @skill + def emergency_land(self) -> str: + """Abort OFFBOARD and land now.""" + conn = self._require_conn() + if conn is None: + return f"{self.drone_key}: NOT CONNECTED" + conn.stop_offboard() + return f"{self.drone_key}: emergency land {'OK' if conn.land() else 'FAILED'}" + + # ------------------------------------------------------------------ + # Vision follow (namespace-local cmd_vel) + # ------------------------------------------------------------------ + + @skill + def start_follow(self, lock_altitude: bool = True) -> str: + """Allow this drone's vision tracker to drive it (OFFBOARD body-frame). + + Safety gate: until this is called, cmd_vel from the tracker is ignored. + The drone should already be armed and airborne. + """ + conn = self._require_conn() + if conn is None: + return f"{self.drone_key}: NOT CONNECTED" + if not conn._offboard_running and not conn.start_offboard(): + return f"{self.drone_key}: failed to enter OFFBOARD" + self._follow_lock_altitude = lock_altitude + self._follow_enabled = True + return f"{self.drone_key}: follow ENABLED (lock_altitude={lock_altitude})" + + @skill + def stop_follow(self) -> str: + """Stop honoring tracker velocities and hold position.""" + self._follow_enabled = False + conn = self._require_conn() + if conn is not None: + try: + conn.set_velocity_body(0.0, 0.0, 0.0, 0.0) + except Exception as e: + logger.debug(f"[{self.drone_key}] zero-velocity on stop_follow failed: {e}") + return f"{self.drone_key}: follow DISABLED" + + def _on_cmd_vel(self, twist: Twist) -> None: + """Forward a tracking Twist — only when follow is explicitly armed.""" + if not self._follow_enabled: + return + conn = self.connection + if conn is None: + return + try: + conn.move_twist(twist, duration=0.0, lock_altitude=self._follow_lock_altitude) + except Exception as e: + logger.warning(f"[{self.drone_key}] cmd_vel forward failed: {e}") + + # ------------------------------------------------------------------ + # Swarm command bus + # ------------------------------------------------------------------ + + def _on_swarm_cmd(self, msg: String) -> None: + """Act on a coordinator broadcast addressed to this drone (or to ``all``).""" + try: + cmd = json.loads(msg.data if hasattr(msg, "data") else str(msg)) + except (ValueError, AttributeError) as e: + logger.warning(f"[{self.drone_key}] malformed swarm_cmd: {e}") + return + + target = cmd.get("target", "all") + if target != "all" and target != self.drone_key: + return + + action = cmd.get("action", "") + args = cmd.get("args", {}) or {} + try: + result = self._dispatch(action, args) + except Exception as e: + logger.warning(f"[{self.drone_key}] swarm_cmd {action} failed: {e}") + return + logger.info(f"[{self.drone_key}] swarm_cmd {action}: {result}") + + def _dispatch(self, action: str, args: dict[str, Any]) -> str: + """Map a swarm-bus action name onto this drone's implementation.""" + if action == "takeoff": + return self.takeoff(float(args.get("altitude", 3.0))) + if action == "land": + return self.land() + if action == "rtl": + return self.rtl() + if action == "hold": + return self.hold() + if action == "arm": + return self.arm() + if action == "kill": + return self.kill() + if action == "emergency_land": + return self.emergency_land() + if action == "goto_ned": + return self._goto_ned( + float(args["north"]), + float(args["east"]), + float(args["down"]), + float(args.get("yaw_rad", 0.0)), + ) + if action == "path": + return self._follow_path( + args.get("path", []), + float(args.get("arrival_radius_m", 2.0)), + ) + return f"unknown action {action!r}" + + +__all__ = ["STATE_PUBLISH_HZ", "Px4DroneConfig", "Px4DroneModule"] diff --git a/dimos/robot/drone/px4_geo.py b/dimos/robot/drone/px4_geo.py new file mode 100644 index 0000000000..c6ecdd7730 --- /dev/null +++ b/dimos/robot/drone/px4_geo.py @@ -0,0 +1,87 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""World-frame geometry helpers shared by the PX4 drone and swarm modules. + +LOCAL_POSITION_NED is referenced to *each drone's own home*, so subtracting two +drones' NED gives nonsense (each reads "0,0,0 at my spawn"). Real inter-drone +distance has to come from GLOBAL_POSITION_INT (lat/lon/alt) and a great-circle +computation — that is what lives here. +""" + +from __future__ import annotations + +import math + +EARTH_RADIUS_M = 6371000.0 +WGS84_EQUATORIAL_M = 6378137.0 + +# A world-frame point: (latitude_deg, longitude_deg, altitude_m). +GlobalPoint = tuple[float, float, float] + + +def haversine_m(lat1: float, lon1: float, lat2: float, lon2: float) -> float: + """Great-circle distance in meters between two (lat, lon) points in degrees.""" + p1 = math.radians(lat1) + p2 = math.radians(lat2) + dp = math.radians(lat2 - lat1) + dl = math.radians(lon2 - lon1) + a = math.sin(dp / 2) ** 2 + math.cos(p1) * math.cos(p2) * math.sin(dl / 2) ** 2 + c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a)) + return EARTH_RADIUS_M * c + + +def offset_latlon(lat: float, lon: float, dn_m: float, de_m: float) -> tuple[float, float]: + """Offset a lat/lon by (north, east) meters. + + Small-angle approximation — accurate to a few centimeters for the sub-km + offsets the swarm deals with. + """ + new_lat = lat + math.degrees(dn_m / WGS84_EQUATORIAL_M) + new_lon = lon + math.degrees(de_m / (WGS84_EQUATORIAL_M * math.cos(math.radians(lat)))) + return new_lat, new_lon + + +def distance_3d_m(g1: GlobalPoint, g2: GlobalPoint) -> float: + """3-D distance between two (lat_deg, lon_deg, alt_m) world-frame points.""" + horiz = haversine_m(g1[0], g1[1], g2[0], g2[1]) + vert = g1[2] - g2[2] + return math.sqrt(horiz * horiz + vert * vert) + + +def pairwise_distances( + positions: dict[str, GlobalPoint], +) -> list[tuple[str, str, float]]: + """Pairwise 3-D distance (meters) between drones' world-frame positions. + + Positions must be (lat_deg, lon_deg, alt_m). Returns ``(a, b, distance)`` + triples with ``a < b`` by key, sorted for stable reporting. + """ + keys = sorted(positions.keys()) + out: list[tuple[str, str, float]] = [] + for i, a in enumerate(keys): + for b in keys[i + 1 :]: + out.append((a, b, distance_3d_m(positions[a], positions[b]))) + return out + + +__all__ = [ + "EARTH_RADIUS_M", + "WGS84_EQUATORIAL_M", + "GlobalPoint", + "distance_3d_m", + "haversine_m", + "offset_latlon", + "pairwise_distances", +] diff --git a/dimos/robot/drone/px4_sitl_connection.py b/dimos/robot/drone/px4_sitl_connection.py new file mode 100644 index 0000000000..52ccf3f65a --- /dev/null +++ b/dimos/robot/drone/px4_sitl_connection.py @@ -0,0 +1,732 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""PX4 SITL flavor of MavlinkConnection. + +PX4 encodes flight modes differently from ArduPilot, so the existing +``MavlinkConnection.set_mode()`` (ArduPilot custom mode IDs) does not work. +This subclass overrides ``set_mode``, ``takeoff``, ``land``, and ``rtl`` for +PX4 semantics, plus adds an OFFBOARD setpoint streamer required for any +external position/velocity control. + +Typical SITL endpoints (from PX4 ``px4-rc.mavlink``): + instance N → ``udp:127.0.0.1:1454N`` (offboard SDK) + sysid = N + 1 +""" + +from __future__ import annotations + +import os +import threading +import time + +from pymavlink import mavutil # type: ignore[import-not-found, import-untyped] + +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.robot.drone.mavlink_connection import MavlinkConnection +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +def _env_float(name: str, default: float) -> float: + raw = os.getenv(name) + if raw is None: + return default + try: + return float(raw) + except ValueError: + logger.warning(f"Invalid {name}={raw!r}; using default {default}") + return default + + +# Vision-tracking safety knobs. Env-overridable so the drone can be tuned in the +# field without code edits (sign flips, speed caps, watchdog timeout). +# * WATCHDOG_S: if no fresh body-velocity setpoint arrives within this window +# (tracker died / video froze), the OFFBOARD streamer commands zero velocity +# (hover-hold) instead of repeating the last command — prevents runaway. +TRACK_BODY_SETPOINT_TIMEOUT_S = _env_float("DIMOS_DRONE_TRACK_WATCHDOG_S", 0.5) +TRACK_MAX_HORIZONTAL_V = _env_float("DIMOS_DRONE_TRACK_MAX_V", 3.0) # m/s, fwd & strafe +TRACK_MAX_VERTICAL_V = _env_float("DIMOS_DRONE_TRACK_MAX_VZ", 1.0) # m/s +TRACK_MAX_YAW_RATE = _env_float("DIMOS_DRONE_TRACK_MAX_YAW", 1.0) # rad/s +TRACK_INVERT_YAW = os.getenv("DIMOS_DRONE_TRACK_INVERT_YAW", "0") == "1" +TRACK_INVERT_LATERAL = os.getenv("DIMOS_DRONE_TRACK_INVERT_LATERAL", "0") == "1" + + +# PX4 main modes (commander/px4_custom_mode.h) +PX4_CUSTOM_MAIN_MODE_MANUAL = 1 +PX4_CUSTOM_MAIN_MODE_ALTCTL = 2 +PX4_CUSTOM_MAIN_MODE_POSCTL = 3 +PX4_CUSTOM_MAIN_MODE_AUTO = 4 +PX4_CUSTOM_MAIN_MODE_ACRO = 5 +PX4_CUSTOM_MAIN_MODE_OFFBOARD = 6 +PX4_CUSTOM_MAIN_MODE_STABILIZED = 7 + +# PX4 AUTO sub-modes +PX4_CUSTOM_SUB_MODE_AUTO_READY = 1 +PX4_CUSTOM_SUB_MODE_AUTO_TAKEOFF = 2 +PX4_CUSTOM_SUB_MODE_AUTO_LOITER = 3 +PX4_CUSTOM_SUB_MODE_AUTO_MISSION = 4 +PX4_CUSTOM_SUB_MODE_AUTO_RTL = 5 +PX4_CUSTOM_SUB_MODE_AUTO_LAND = 6 + +# Friendly mode name → (main, sub) +PX4_MODE_TABLE: dict[str, tuple[int, int]] = { + "MANUAL": (PX4_CUSTOM_MAIN_MODE_MANUAL, 0), + "POSCTL": (PX4_CUSTOM_MAIN_MODE_POSCTL, 0), + "ALTCTL": (PX4_CUSTOM_MAIN_MODE_ALTCTL, 0), + "OFFBOARD": (PX4_CUSTOM_MAIN_MODE_OFFBOARD, 0), + "STABILIZED": (PX4_CUSTOM_MAIN_MODE_STABILIZED, 0), + "AUTO.LOITER": (PX4_CUSTOM_MAIN_MODE_AUTO, PX4_CUSTOM_SUB_MODE_AUTO_LOITER), + "AUTO.TAKEOFF": (PX4_CUSTOM_MAIN_MODE_AUTO, PX4_CUSTOM_SUB_MODE_AUTO_TAKEOFF), + "AUTO.LAND": (PX4_CUSTOM_MAIN_MODE_AUTO, PX4_CUSTOM_SUB_MODE_AUTO_LAND), + "AUTO.RTL": (PX4_CUSTOM_MAIN_MODE_AUTO, PX4_CUSTOM_SUB_MODE_AUTO_RTL), + "AUTO.MISSION": (PX4_CUSTOM_MAIN_MODE_AUTO, PX4_CUSTOM_SUB_MODE_AUTO_MISSION), +} + + +def default_offboard_endpoint(instance: int = 0, host: str = "127.0.0.1") -> str: + """Return the PX4 SITL offboard MAVLink endpoint for a given instance.""" + return f"udp:{host}:{14540 + instance}" + + +class Px4SitlConnection(MavlinkConnection): + """MAVLink connection tuned for PX4 SITL. + + Adds: + * PX4 mode encoding (main_mode/sub_mode in MAV_CMD_DO_SET_MODE). + * OFFBOARD setpoint streamer that keeps PX4 in OFFBOARD by re-sending + the latest setpoint at ≥ 2 Hz (PX4 falls out of OFFBOARD otherwise). + * RTL helper. + """ + + OFFBOARD_STREAM_HZ = 20.0 # comfortably above PX4's 2 Hz minimum + + def __init__( + self, + connection_string: str | None = None, + instance: int = 0, + outdoor: bool = False, + max_velocity: float = 5.0, + ) -> None: + if connection_string is None: + connection_string = default_offboard_endpoint(instance) + # PX4 SITL maps instance N to MAV_SYS_ID = N + 1. Pass it down so + # MavlinkConnection.connect() can skip HEARTBEAT discovery on instances + # >= 1 (which don't stream HEARTBEAT on their Onboard channel reliably). + super().__init__( + connection_string=connection_string, + outdoor=outdoor, + max_velocity=max_velocity, + target_system=instance + 1, + # GCS presence is fleet-level, not per-drone: every PX4 instance + # expects its ground station at the ONE shared port 14550, so + # SwarmCoordinator._gcs_presence_loop owns that socket. Per-drone + # attempts (pushing at 18570+i, or binding 14550+i) are both dead + # ends; see the notes on that loop. + ) + self.instance = instance + self._offboard_setpoint: tuple[float, float, float, float] | None = None # (vn, ve, vd, yaw_rate) + # Body-frame velocity setpoint (vx_fwd, vy_right, vz_down, yaw_rate). Used by + # vision tracking, which thinks in "forward/strafe relative to the nose" + yaw. + self._offboard_body_setpoint: tuple[float, float, float, float] | None = None + # Monotonic timestamp of the last body setpoint update, for the watchdog. + self._body_setpoint_ts: float = 0.0 + self._offboard_position: tuple[float, float, float, float] | None = None # (n, e, d, yaw) + # Path mode: a list of (n, e, d, yaw) waypoints, plus current index + arrival radius. + # When _offboard_path is non-empty, the streamer walks the list, advancing once + # the drone is within _path_arrival_radius_m of the current target. After the + # last waypoint the drone holds in place (path becomes a sticky setpoint). + self._offboard_path: list[tuple[float, float, float, float]] | None = None + self._path_index: int = 0 + self._path_arrival_radius_m: float = 2.0 + self._offboard_thread: threading.Thread | None = None + self._offboard_running = False + self._offboard_lock = threading.Lock() + + # ------------------------------------------------------------------ + # Mode handling (PX4-specific) + # ------------------------------------------------------------------ + + def set_mode(self, mode: str) -> bool: + """Set PX4 flight mode by friendly name (e.g. ``OFFBOARD``, ``AUTO.RTL``).""" + if not self.connected: + return False + if mode not in PX4_MODE_TABLE: + logger.error(f"Unknown PX4 mode: {mode}. Valid: {sorted(PX4_MODE_TABLE)}") + return False + main, sub = PX4_MODE_TABLE[mode] + return self._set_mode_px4(main, sub, label=mode) + + def _set_mode_px4(self, main: int, sub: int, label: str = "") -> bool: + logger.info(f"PX4 set_mode → {label or f'main={main} sub={sub}'}") + self.mavlink.mav.command_long_send( + self.mavlink.target_system, + self.mavlink.target_component, + mavutil.mavlink.MAV_CMD_DO_SET_MODE, + 0, + mavutil.mavlink.MAV_MODE_FLAG_CUSTOM_MODE_ENABLED, + float(main), + float(sub), + 0, + 0, + 0, + 0, + ) + ack = self.mavlink.recv_match(type="COMMAND_ACK", blocking=True, timeout=3) + if ack and ack.result == mavutil.mavlink.MAV_RESULT_ACCEPTED: + return True + logger.warning(f"set_mode {label} not accepted (ack={ack})") + return False + + # ------------------------------------------------------------------ + # Auto skills (mode-driven; no setpoint stream required) + # ------------------------------------------------------------------ + + def takeoff(self, altitude: float = 3.0) -> bool: + """Arm and switch to AUTO.TAKEOFF at the requested relative altitude.""" + if not self.connected: + return False + # Any leftover OFFBOARD setpoint stream from a previous skill (grid_sweep, + # goto_drone, etc.) must be torn down BEFORE we re-arm, or PX4 will keep + # chasing stale waypoints the moment it accepts the new mode. + self.stop_offboard() + # Clear whatever nav state the last command left behind. After an RTL + # completes and the vehicle auto-disarms, PX4 stays latched in AUTO.RTL + # and denies the next arm with "Resolve system health failures first" -- + # with no Preflight Fail line to explain it. Dropping to AUTO.LOITER + # first makes the obvious demo loop (take off, RTL, take off again) + # work. Best effort: a failed mode switch is not itself fatal. + if not self.set_mode("AUTO.LOITER"): + logger.debug("takeoff: could not pre-set AUTO.LOITER; arming anyway") + if not self.arm(): + logger.error("PX4 takeoff failed: arm rejected") + return False + # MAV_CMD_NAV_TAKEOFF param7 is altitude ABOVE MEAN SEA LEVEL, not above + # home. Passing a relative altitude makes PX4 compare, say, 5 m AMSL + # against a vehicle already sitting at 488 m and reply "Already higher + # than takeoff altitude" -- it silently does nothing. Convert using the + # home altitude implied by telemetry: home_amsl = alt - relative_alt. + target_amsl = self._relative_to_amsl(altitude) + if target_amsl is None: + logger.error( + "PX4 takeoff failed: no GLOBAL_POSITION_INT yet, cannot convert " + f"{altitude}m relative to an absolute altitude" + ) + return False + self.mavlink.mav.command_long_send( + self.mavlink.target_system, + self.mavlink.target_component, + mavutil.mavlink.MAV_CMD_NAV_TAKEOFF, + 0, + 0, + 0, + 0, + float("nan"), # yaw → leave unchanged + float("nan"), # lat → use current + float("nan"), # lon → use current + target_amsl, + ) + logger.info( + f"PX4 takeoff command sent (altitude={altitude}m relative " + f"-> {target_amsl:.1f}m AMSL)" + ) + return True + + def _relative_to_amsl(self, relative_alt_m: float) -> float | None: + """Convert an altitude above home into altitude above mean sea level. + + Returns None when no GLOBAL_POSITION_INT has arrived yet, since guessing + an absolute altitude would send the vehicle somewhere unintended. + """ + gpi = self.telemetry.get("GLOBAL_POSITION_INT") + if not gpi: + return None + amsl = gpi.get("alt") + rel = gpi.get("relative_alt") + if amsl is None or rel is None: + return None + return float(amsl) - float(rel) + relative_alt_m + + def land(self) -> bool: + """Switch to AUTO.LAND.""" + if not self.connected: + return False + # Same reasoning as takeoff: stop the OFFBOARD streamer so PX4 isn't + # being pulled in two directions while it tries to land. + self.stop_offboard() + return self.set_mode("AUTO.LAND") + + def rtl(self) -> bool: + """Switch to AUTO.RTL.""" + if not self.connected: + return False + self.stop_offboard() + return self.set_mode("AUTO.RTL") + + def hold(self) -> bool: + """Switch to AUTO.LOITER (hover in place).""" + if not self.connected: + return False + self.stop_offboard() + return self.set_mode("AUTO.LOITER") + + # ------------------------------------------------------------------ + # OFFBOARD setpoint streaming + # ------------------------------------------------------------------ + + def start_offboard( + self, + initial_setpoint: tuple[float, float, float, float] | None = None, + ) -> bool: + """Begin streaming OFFBOARD setpoints and switch PX4 into OFFBOARD mode. + + Args: + initial_setpoint: (vn, ve, vd, yaw_rate). Defaults to hover (all zeros). + + Returns: + True if PX4 accepted the OFFBOARD mode switch. + + Failure semantics: on ``set_mode("OFFBOARD")`` rejection we **tear the + streamer back down** so a subsequent call retries from a clean slate. + Previously a failed mode switch left ``_offboard_running=True`` with a + live thread but PX4 not in OFFBOARD — later callers checked + ``if not _offboard_running`` and skipped the retry, leaving the drone + stranded. + """ + if self._offboard_running: + return True + + with self._offboard_lock: + self._offboard_setpoint = initial_setpoint or (0.0, 0.0, 0.0, 0.0) + self._offboard_body_setpoint = None + self._offboard_position = None + self._offboard_path = None + self._offboard_running = True + self._offboard_thread = threading.Thread( + target=self._offboard_loop, daemon=True, name=f"px4-offboard-{self.instance}" + ) + self._offboard_thread.start() + + # PX4 needs to see ≥ 1 second of setpoints before it accepts OFFBOARD. + time.sleep(1.1) + if self.set_mode("OFFBOARD"): + return True + + # Mode switch rejected — undo so the next call retries cleanly. + logger.warning( + "start_offboard: PX4 rejected mode switch, tearing streamer down for retry" + ) + self.stop_offboard() + return False + + def stop_offboard(self) -> None: + """Halt the OFFBOARD setpoint streamer and clear any pending setpoints. + + Safe to call when the streamer isn't running. Always clears + ``_offboard_path``, ``_offboard_position``, and ``_offboard_setpoint`` + so a subsequent ``start_offboard()`` starts from a clean slate (no stale + waypoints leaking between commands). + """ + self._offboard_running = False + if self._offboard_thread and self._offboard_thread.is_alive(): + self._offboard_thread.join(timeout=1.0) + self._offboard_thread = None + with self._offboard_lock: + self._offboard_path = None + self._path_index = 0 + self._offboard_position = None + self._offboard_setpoint = None + + def set_velocity_ned( + self, vn: float, ve: float, vd: float, yaw_rate: float = 0.0 + ) -> None: + """Update the streamed OFFBOARD velocity setpoint (NED frame, m/s). + + Cancels any active position setpoint AND any active path — the streamer + will switch to pure velocity mode on the next tick. + """ + with self._offboard_lock: + self._offboard_setpoint = (vn, ve, vd, yaw_rate) + self._offboard_body_setpoint = None + self._offboard_position = None + self._offboard_path = None + + def set_position_ned( + self, n: float, e: float, d: float, yaw: float = 0.0 + ) -> None: + """Update the streamed OFFBOARD position setpoint (local NED, meters). + + ``d`` is positive-down (NED). To fly to 5m altitude, pass ``d = -5``. + """ + with self._offboard_lock: + self._offboard_position = (n, e, d, yaw) + self._offboard_setpoint = None + self._offboard_body_setpoint = None + self._offboard_path = None # explicit position cancels any active path + + def set_position_path( + self, + waypoints: list[tuple[float, float, float, float]], + arrival_radius_m: float = 2.0, + ) -> None: + """Stream a sequence of OFFBOARD position waypoints. + + Each waypoint is (n, e, d, yaw) in local NED. ``d`` is positive-down. + The streamer holds each waypoint until the drone is within + ``arrival_radius_m`` of it, then advances. After the final waypoint + the drone hovers at that point (the path becomes a sticky setpoint). + + This is the foundation for multi-leg patterns like a boustrophedon + ("lawnmower") sweep — the fleet module builds the list, this method + executes it. + + Args: + waypoints: ordered list of (north_m, east_m, down_m, yaw_rad). + Must contain at least one element. + arrival_radius_m: 3-D distance (meters) at which a waypoint counts + as "reached". Default 2 m, which matches PX4 SITL position + tracking accuracy at typical cruise speeds. + + Raises: + ValueError: if ``waypoints`` is empty. + """ + if not waypoints: + raise ValueError("set_position_path requires at least one waypoint") + with self._offboard_lock: + self._offboard_path = list(waypoints) + self._path_index = 0 + self._path_arrival_radius_m = float(arrival_radius_m) + # Path mode wins over the single-point modes. + self._offboard_position = None + self._offboard_setpoint = None + self._offboard_body_setpoint = None + + def _path_advance_if_arrived(self, target: tuple[float, float, float, float]) -> None: + """If the drone is within arrival radius of ``target``, advance the path index. + + Holds at the last waypoint (no wraparound). + """ + if self._offboard_path is None: + return + cur = self.get_local_ned() + if cur is None: + return + dn = cur[0] - target[0] + de = cur[1] - target[1] + dd = cur[2] - target[2] + dist = (dn * dn + de * de + dd * dd) ** 0.5 + if dist <= self._path_arrival_radius_m: + with self._offboard_lock: + if ( + self._offboard_path is not None + and self._path_index < len(self._offboard_path) - 1 + ): + self._path_index += 1 + + def _offboard_loop(self) -> None: + period = 1.0 / self.OFFBOARD_STREAM_HZ + while self._offboard_running: + try: + with self._offboard_lock: + path = self._offboard_path + idx = self._path_index + pos = self._offboard_position + vel = self._offboard_setpoint + body_vel = self._offboard_body_setpoint + if path is not None: + target = path[min(idx, len(path) - 1)] + self._send_position_target(*target) + self._path_advance_if_arrived(target) + elif pos is not None: + self._send_position_target(*pos) + elif body_vel is not None: + send = self._body_setpoint_to_send(time.monotonic()) + if send is not None: + self._send_velocity_target_body(*send) + else: + vn, ve, vd, yaw_rate = vel or (0.0, 0.0, 0.0, 0.0) + self._send_velocity_target(vn, ve, vd, yaw_rate) + except Exception as e: # never let the streamer thread die silently + logger.debug(f"offboard stream send error: {e}") + time.sleep(period) + + def _send_velocity_target(self, vn: float, ve: float, vd: float, yaw_rate: float) -> None: + # type_mask: ignore position, accel, yaw — use velocity + yaw_rate. + # Bit layout (MAVLink spec): pos(0-2) vel(3-5) acc(6-8) force(9) yaw(10) yaw_rate(11) + type_mask = 0b0000_1011_1100_0111 + self.mavlink.mav.set_position_target_local_ned_send( + 0, + self.mavlink.target_system, + self.mavlink.target_component, + mavutil.mavlink.MAV_FRAME_LOCAL_NED, + type_mask, + 0, 0, 0, + float(vn), float(ve), float(vd), + 0, 0, 0, + 0.0, + float(yaw_rate), + ) + + def _send_velocity_target_body( + self, vx: float, vy: float, vz: float, yaw_rate: float + ) -> None: + """Stream a BODY-frame velocity setpoint (vx fwd, vy right, vz down). + + Unlike the world-NED ``_send_velocity_target``, this uses + ``MAV_FRAME_BODY_NED`` so the autopilot rotates the command by the + drone's current heading — "go forward / strafe relative to the nose". + The type_mask uses velocity + yaw_rate (ignores position, accel, and the + absolute yaw angle) so vision tracking can rotate the drone to keep the + target centred. + """ + # Bit layout (set = ignore): pos(0-2) vel(3-5) acc(6-8) force(9) yaw(10) yaw_rate(11) + # Use vel + yaw_rate -> ignore pos, acc, force, and absolute yaw. + type_mask = 0b0000_0111_1100_0111 + self.mavlink.mav.set_position_target_local_ned_send( + 0, + self.mavlink.target_system, + self.mavlink.target_component, + mavutil.mavlink.MAV_FRAME_BODY_NED, + type_mask, + 0, 0, 0, + float(vx), float(vy), float(vz), + 0, 0, 0, + 0.0, + float(yaw_rate), + ) + + def set_velocity_body( + self, vx: float, vy: float, vz: float, yaw_rate: float = 0.0 + ) -> None: + """Update the streamed OFFBOARD body-frame velocity setpoint (m/s, rad/s). + + vx = forward, vy = right, vz = down (positive). Cancels any active + world-velocity / position / path setpoint. + """ + with self._offboard_lock: + self._offboard_body_setpoint = (vx, vy, vz, yaw_rate) + self._body_setpoint_ts = time.monotonic() + self._offboard_setpoint = None + self._offboard_position = None + self._offboard_path = None + + def _body_setpoint_to_send( + self, now: float + ) -> tuple[float, float, float, float] | None: + """Resolve the body setpoint to stream, applying the staleness watchdog. + + Returns None when no body setpoint is active. Returns a zero setpoint + (hover-hold) when the latest setpoint is older than the watchdog window + — this is the failsafe that stops the drone if the tracker or video feed + dies mid-flight instead of repeating a stale "keep moving" command. + """ + with self._offboard_lock: + sp = self._offboard_body_setpoint + ts = self._body_setpoint_ts + if sp is None: + return None + if (now - ts) > TRACK_BODY_SETPOINT_TIMEOUT_S: + return (0.0, 0.0, 0.0, 0.0) + return sp + + @staticmethod + def _shape_track_twist( + twist: Twist, lock_altitude: bool + ) -> tuple[float, float, float, float]: + """Convert a tracking Twist into a safe, sign-corrected body setpoint. + + Applies the env-tunable sign inversions and velocity/yaw caps. Pure + function (no I/O) so it is unit-testable without a live MAVLink link. + """ + vx = float(twist.linear.x) # forward + vy = float(twist.linear.y) # right + if TRACK_INVERT_LATERAL: + vy = -vy + vz = 0.0 if lock_altitude else -float(twist.linear.z) + yaw_rate = float(twist.angular.z) + if TRACK_INVERT_YAW: + yaw_rate = -yaw_rate + + vx = max(-TRACK_MAX_HORIZONTAL_V, min(TRACK_MAX_HORIZONTAL_V, vx)) + vy = max(-TRACK_MAX_HORIZONTAL_V, min(TRACK_MAX_HORIZONTAL_V, vy)) + vz = max(-TRACK_MAX_VERTICAL_V, min(TRACK_MAX_VERTICAL_V, vz)) + yaw_rate = max(-TRACK_MAX_YAW_RATE, min(TRACK_MAX_YAW_RATE, yaw_rate)) + return vx, vy, vz, yaw_rate + + def move_twist( + self, twist: Twist, duration: float = 0.0, lock_altitude: bool = True + ) -> None: + """Drive the drone from a ROS-style Twist via the OFFBOARD streamer. + + Mirrors the Tello/MAVLink ``move_twist`` contract so the same + ``DroneTrackingModule`` cmd_vel output works on PX4: + linear.x -> forward (body) + linear.y -> right (body) + linear.z -> up (used only when lock_altitude is False) + angular.z -> yaw rate (rad/s) [honoured here, unlike mavlink_connection] + + Sign inversions and speed caps are applied (env-tunable) for safe field + bring-up, and the OFFBOARD streamer's watchdog will zero the command if + updates stop arriving. + """ + vx, vy, vz, yaw_rate = self._shape_track_twist(twist, lock_altitude) + + if not self._offboard_running: + self.start_offboard((0.0, 0.0, 0.0, 0.0)) + self.set_velocity_body(vx, vy, vz, yaw_rate) + + if duration > 0: + time.sleep(duration) + self.set_velocity_body(0.0, 0.0, 0.0, 0.0) + + def _send_position_target(self, n: float, e: float, d: float, yaw: float) -> None: + # type_mask: use position + yaw, ignore velocity, accel, yaw_rate. + type_mask = 0b0000_1011_1111_1000 + self.mavlink.mav.set_position_target_local_ned_send( + 0, + self.mavlink.target_system, + self.mavlink.target_component, + mavutil.mavlink.MAV_FRAME_LOCAL_NED, + type_mask, + float(n), float(e), float(d), + 0, 0, 0, + 0, 0, 0, + float(yaw), + 0.0, + ) + + # ------------------------------------------------------------------ + # Body-frame ROS-style move() override (uses OFFBOARD streamer) + # ------------------------------------------------------------------ + + def move(self, velocity: Vector3, duration: float = 0.0) -> None: + """Send a body-frame velocity command via OFFBOARD. + + ROS convention: x = forward (m/s), y = left, z = up. Translates to + NED setpoints based on the latest yaw from telemetry. + """ + # Convert ROS body to NED. Without yaw rotation, treat x=north, y=-east, z=-down. + # For body-relative motion we'd need yaw; for now keep it simple and use world NED + # which matches how OFFBOARD setpoints are expected. + vn = float(velocity.x) + ve = -float(velocity.y) + vd = -float(velocity.z) + + if not self._offboard_running: + self.start_offboard((vn, ve, vd, 0.0)) + else: + self.set_velocity_ned(vn, ve, vd) + + if duration > 0: + time.sleep(duration) + self.set_velocity_ned(0.0, 0.0, 0.0) + + def disconnect(self) -> None: + self.stop_offboard() + if hasattr(super(), "disconnect"): + super().disconnect() # type: ignore[misc] + else: + self.connected = False + + # ------------------------------------------------------------------ + # Convenience accessors used by fleet aggregation + # ------------------------------------------------------------------ + + def get_local_ned(self) -> tuple[float, float, float] | None: + """Return current LOCAL_POSITION_NED if available.""" + local = self.telemetry.get("LOCAL_POSITION_NED") + if not local: + return None + return ( + float(local.get("x", 0.0)), + float(local.get("y", 0.0)), + float(local.get("z", 0.0)), + ) + + def get_global_position(self) -> tuple[float, float, float] | None: + """Return current world-frame position as (lat_deg, lon_deg, rel_alt_m). + + Returns None if telemetry hasn't produced a GLOBAL_POSITION_INT yet. + + Note: ``MavlinkConnection.update_telemetry`` already normalises lat/lon + to degrees and relative_alt to meters in-place when stamping the + GLOBAL_POSITION_INT message into ``self.telemetry`` (see + ``mavlink_connection.py`` around line 204). So we just read the fields + as-is — no further unit conversion. + + This is the *correct* basis for inter-drone distance: each drone's + LOCAL_POSITION_NED is referenced to its own home, so subtracting two + local NEDs gives nonsense. Subtracting two lat/lons (Haversine) gives + real meters. + """ + gp = self.telemetry.get("GLOBAL_POSITION_INT") + if not gp: + return None + lat = gp.get("lat") + lon = gp.get("lon") + if lat is None or lon is None: + return None + rel_alt = gp.get("relative_alt", 0) or 0 + return (float(lat), float(lon), float(rel_alt)) + + def get_battery_pct(self) -> float | None: + sys_status = self.telemetry.get("SYS_STATUS") + if not sys_status: + return None + pct = sys_status.get("battery_remaining") + if pct is None or pct < 0: + return None + return float(pct) + + def get_armed(self) -> bool: + hb = self.telemetry.get("HEARTBEAT") + if not hb: + return False + return bool(hb.get("armed", False)) + + def get_mode(self) -> str | None: + """Current PX4 flight mode as a friendly name, e.g. ``AUTO.RTL``. + + Without this an operator cannot tell OFFBOARD from AUTO.RTL in fleet + telemetry, which is what made a latched AUTO.RTL (a drone that silently + refuses to arm again) so hard to diagnose: every other field looked + healthy. + + PX4 packs the mode into HEARTBEAT.custom_mode as + ``main = (custom_mode >> 16) & 0xFF`` and ``sub = (custom_mode >> 24) & 0xFF``. + """ + hb = self.telemetry.get("HEARTBEAT") + if not hb: + return None + custom = hb.get("custom_mode") + if custom is None or custom < 0: + return None + main = (int(custom) >> 16) & 0xFF + sub = (int(custom) >> 24) & 0xFF + for name, (m, sm) in PX4_MODE_TABLE.items(): + if m == main and sm == sub: + return name + # Unmapped combinations are real (e.g. AUTO.MISSION); report the raw + # numbers rather than None so the operator still sees a change. + return f"UNKNOWN({main}.{sub})" + + +__all__ = [ + "PX4_MODE_TABLE", + "Px4SitlConnection", + "default_offboard_endpoint", +] diff --git a/dimos/robot/drone/px4_sitl_fleet_config.py b/dimos/robot/drone/px4_sitl_fleet_config.py new file mode 100644 index 0000000000..903fd11d21 --- /dev/null +++ b/dimos/robot/drone/px4_sitl_fleet_config.py @@ -0,0 +1,119 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fleet configuration for PX4 SITL multi-vehicle simulations. + +The PX4 multi-vehicle convention (and the in-repo ``dimos/simulation/px4_hil`` simulator) +maps SITL instance N to: + * MAVLink offboard endpoint: ``udp:127.0.0.1:1454N`` + * MAVLink GCS port: 14550 + N + * MAV_SYS_ID: N + 1 + * Gazebo model name: ``x500_N`` + +This module turns the DimOS ``--robot-ips`` flag (or a default fleet size of 3, +matching the sim demo deliverable) into a list of ``Px4SitlDroneConfig`` +records that the fleet module can consume. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from dimos.core.global_config import GlobalConfig, global_config + +DEFAULT_SITL_HOST = "127.0.0.1" +DEFAULT_OFFBOARD_PORT_BASE = 14540 +DEFAULT_GCS_PORT_BASE = 14550 +DEFAULT_FLEET_SIZE = 3 # matches the X500 sim demo deliverable + + +@dataclass(frozen=True) +class Px4SitlDroneConfig: + """Per-drone identity for a PX4 SITL instance.""" + + key: str + instance: int + host: str + offboard_port: int + gcs_port: int + sys_id: int + model_name: str + + @property + def connection_string(self) -> str: + return f"udp:{self.host}:{self.offboard_port}" + + +def _split_csv(raw: str | None) -> list[str]: + if raw is None: + return [] + return [item.strip() for item in raw.split(",") if item.strip()] + + +def get_px4_sitl_fleet_configs( + cfg: GlobalConfig = global_config, + fleet_size: int | None = None, +) -> list[Px4SitlDroneConfig]: + """Return the configured PX4 SITL fleet. + + Resolution order: + 1. Explicit ``fleet_size`` argument (used by blueprints to pick a default). + 2. CSV count from ``--robot-ips`` (each entry overrides the host). + 3. ``DEFAULT_FLEET_SIZE``. + """ + hosts = _split_csv(cfg.robot_ips) + if fleet_size is None: + fleet_size = len(hosts) if hosts else DEFAULT_FLEET_SIZE + if fleet_size < 1: + raise ValueError(f"fleet_size must be >= 1, got {fleet_size}") + + drones: list[Px4SitlDroneConfig] = [] + for index in range(fleet_size): + host = hosts[index] if index < len(hosts) else (hosts[0] if hosts else DEFAULT_SITL_HOST) + drones.append( + Px4SitlDroneConfig( + key=f"drone-{index + 1}", + instance=index, + host=host, + offboard_port=DEFAULT_OFFBOARD_PORT_BASE + index, + gcs_port=DEFAULT_GCS_PORT_BASE + index, + sys_id=index + 1, + model_name=f"x500_{index}", + ) + ) + return drones + + +def format_px4_sitl_fleet_prompt_block( + cfg: GlobalConfig = global_config, + fleet_size: int | None = None, +) -> str: + """Render a concise prompt section describing the configured fleet.""" + drones = get_px4_sitl_fleet_configs(cfg, fleet_size=fleet_size) + return "\n".join( + f"- {d.key}: instance={d.instance}, sys_id={d.sys_id}, " + f"model={d.model_name}, offboard={d.connection_string}" + for d in drones + ) + + +__all__ = [ + "DEFAULT_FLEET_SIZE", + "DEFAULT_GCS_PORT_BASE", + "DEFAULT_OFFBOARD_PORT_BASE", + "DEFAULT_SITL_HOST", + "Px4SitlDroneConfig", + "format_px4_sitl_fleet_prompt_block", + "get_px4_sitl_fleet_configs", +] diff --git a/dimos/robot/drone/px4_swarm_coordinator.py b/dimos/robot/drone/px4_swarm_coordinator.py new file mode 100644 index 0000000000..9dcddad314 --- /dev/null +++ b/dimos/robot/drone/px4_swarm_coordinator.py @@ -0,0 +1,1335 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fleet-level brain for the PX4 swarm: one shared instance, N namespaced drones. + +This is the other half of the ``Px4SitlFleetModule`` split. Where +[Px4DroneModule][dimos.robot.drone.px4_drone_module.Px4DroneModule] owns one +vehicle inside its own namespace, this module sits *outside* all namespaces and +owns everything that is inherently about the fleet as a whole: + +* the shared world-frame picture (``fleet_state``, ``count_within``), +* the minimum-separation guardrail, which is only meaningful across drones, +* the multi-drone demo maneuvers (``grid_sweep``, ``line_formation``, + ``investigate``), and +* the aggregate commands (``takeoff_all``, ``rtl_all``, ``kill_all``, ...). + +It never holds a MAVLink connection. Data crosses the namespace boundary on two +exposed streams: it *listens* on ``drone_state`` (every drone publishes its own +snapshot there) and *broadcasts* on ``swarm_cmd`` (each drone acts only on +messages addressed to its key, or to ``all``). + +That indirection is what makes the fleet size dynamic: adding a fourth drone is +one more namespaced blueprint, with no change here. +""" + +from __future__ import annotations + +import json +import math +import os +import threading +import time +from typing import Any + +from dimos_lcm.std_msgs import String + +from dimos.agents.annotation import skill +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.robot.drone.px4_geo import ( + GlobalPoint, + distance_3d_m, + offset_latlon, + pairwise_distances, +) +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +# A drone whose last snapshot is older than this is reported as stale rather +# than silently trusted for guardrail math. +STATE_STALE_AFTER_SEC = 5.0 + +# Separate, much tighter bound for the separation guardrail. Reporting a robot +# as "present" can tolerate seconds of lag; deciding two vehicles will not +# collide cannot. At a 5 m/s closing speed, 5 s of staleness is 25 m of travel +# against a 2 m floor -- the check would be arithmetic theatre. The guardrail +# refuses to answer on data older than this rather than answering wrongly. +GUARDRAIL_MAX_AGE_SEC = 0.75 + + +class SwarmCoordinatorConfig(ModuleConfig): + """Fleet-wide policy. Not per-drone — these apply across the whole swarm.""" + + # Minimum allowed world-frame separation between any two drones, in meters. + # Position commands predicted to violate this are rejected. + min_separation_m: float = 2.0 + # Optional fleet-wide altitude ceiling (m above home, positive up). Each + # drone also enforces its own cap; this one catches fleet maneuvers early. + max_altitude_m: float | None = None + # Expected drone keys, CSV (e.g. "drone1,drone2,drone3"). Used only to + # report drones that have never reported in. Empty = infer from traffic. + expected_drones: str = "" + # Battery floor for preflight_check, in percent. PX4 has its own low-battery + # failsafe; this is the "don't even take off" gate above it. + min_battery_pct: float = 25.0 + # UDP port where PX4 expects its ground control station. Every PX4 + # instance unicasts its GCS-mode MAVLink stream at this ONE port (14550, + # the QGroundControl convention -- verified: instances 0 and 1 both target + # 14550), and only credits "connected to ground control station" for + # heartbeats sourced FROM it. That check becomes an arming blocker once the + # datalink-loss failsafe (NAV_DLL_ACT) is configured, so the coordinator -- + # one per fleet, like a real GCS -- owns this socket and answers every + # vehicle from it. When DimOS dies these heartbeats stop and every drone + # returns and lands on its own authority, which is exactly the failsafe + # story wanted on hardware. 0 disables (e.g. when a real QGC is attached). + gcs_port: int = 14550 + # Operating radius: maximum horizontal distance (m) of any commanded + # waypoint from the drone's home. Dispatch-time gate; the PX4-side + # geofence (GF_ACTION=hold at GF_MAX_HOR_DIST, written by + # sim_params.py) is the in-flight backstop for velocity commands and + # drift, which no dispatch check can see. The sim ground plane ends at + # 300 m, hence the default. + max_range_m: float = 250.0 + + +class SwarmCoordinator(Module): + """Fleet state, spacing guardrails, and multi-drone maneuvers.""" + + config: SwarmCoordinatorConfig + + # Exposed streams (global): telemetry in from every drone, commands out. + drone_state: In[String] + swarm_cmd: Out[String] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._states: dict[str, dict[str, Any]] = {} + self._lock = threading.RLock() + self._seq = 0 + # Google Maps client is built lazily on first map skill so a missing + # GOOGLE_MAPS_API_KEY doesn't break startup. None = not tried, + # False = tried and unavailable. + self._gmaps_client: Any = None + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + @rpc + def start(self) -> None: + super().start() + if getattr(self.drone_state, "transport", None): + self.drone_state.subscribe(self._on_drone_state) + logger.info("SwarmCoordinator subscribed to drone_state") + if self.config.gcs_port: + self._gcs_stop = threading.Event() + self._gcs_thread = threading.Thread( + target=self._gcs_presence_loop, name="fleet-gcs-presence", daemon=True + ) + self._gcs_thread.start() + + @rpc + def stop(self) -> None: + ev = getattr(self, "_gcs_stop", None) + if ev is not None: + ev.set() + super().stop() + + def _gcs_presence_loop(self) -> None: + """Be the fleet's ground station, the way QGroundControl is. + + Binds the shared GCS port, discovers every PX4 instance from the + source address of its stream, and answers each with a 2 Hz GCS + heartbeat FROM that port. Hard-won details, learned by strace: + + * Every instance targets the ONE port; a per-drone socket cannot work + (only one process can bind it) and letting pymavlink reply to + "whoever sent last" splits the heartbeat stream across instances -- + each then sees ~0.5 Hz with gaps beyond the 2.5 s heartbeat validity + window, so the GCS flag flaps ("connection regained" once a second) + and arming fails at random. + * Heartbeats PUSHED at an instance's local port from an ephemeral + source are received (they show in `mavlink status` rx counters) but + never credited. Only traffic sourced from the GCS port counts. + * 2 Hz leaves margin under the 2.5 s window even when this thread + jitters under load. + """ + try: + from pymavlink import mavutil # deferred: pymavlink is heavy + except Exception as e: + logger.warning(f"fleet GCS presence disabled: pymavlink unavailable ({e})") + return + try: + link = mavutil.mavlink_connection( + f"udp:0.0.0.0:{self.config.gcs_port}", + source_system=255, + source_component=190, + ) + except Exception as e: + logger.warning( + f"fleet GCS presence disabled: cannot bind :{self.config.gcs_port} ({e}). " + "Is QGroundControl running? Set gcs_port=0 to silence this." + ) + return + logger.info(f"fleet GCS presence up on :{self.config.gcs_port}") + stop = self._gcs_stop + ticks = drained = sent = send_err = 0 + while not stop.is_set(): + ticks += 1 + # Drain a slice of the flood. In pymavlink's server mode this is + # what maintains `link.clients` -- every vehicle whose stream we + # receive is registered there with a freshness stamp. + for _ in range(400): + try: + if link.recv_match(blocking=False) is None: + break + drained += 1 + except Exception: + break + try: + # Server-mode write() fans out to EVERY live client, so one + # heartbeat_send reaches the whole fleet. + link.mav.heartbeat_send( + mavutil.mavlink.MAV_TYPE_GCS, + mavutil.mavlink.MAV_AUTOPILOT_INVALID, + 0, + 0, + 0, + ) + sent += 1 + except Exception as e: + send_err += 1 + if send_err <= 3: + logger.warning(f"fleet GCS heartbeat failed: {type(e).__name__}: {e}") + if ticks % 120 == 0: + clients = sorted(getattr(link, "clients", []) or []) + logger.debug( + f"fleet GCS presence: clients={clients} sent={sent} " + f"send_err={send_err} drained={drained}" + ) + stop.wait(0.5) + try: + link.close() + except Exception: + pass + + def _on_drone_state(self, msg: String) -> None: + try: + state = json.loads(msg.data if hasattr(msg, "data") else str(msg)) + except (ValueError, AttributeError) as e: + logger.debug(f"malformed drone_state: {e}") + return + key = state.get("key") + if not key: + return + with self._lock: + self._states[key] = state + + # ------------------------------------------------------------------ + # Fleet view + # ------------------------------------------------------------------ + + @property + def _expected_keys(self) -> list[str]: + raw = self.config.expected_drones or "" + return [k.strip() for k in raw.split(",") if k.strip()] + + def _snapshot(self) -> dict[str, dict[str, Any]]: + """Copy of the latest per-drone states, with staleness annotated.""" + now = time.time() + with self._lock: + snap = {k: dict(v) for k, v in self._states.items()} + for state in snap.values(): + age = now - float(state.get("ts", 0.0)) + state["age_sec"] = round(age, 2) + state["stale"] = age > STATE_STALE_AFTER_SEC + for key in self._expected_keys: + if key not in snap: + snap[key] = {"key": key, "connected": False, "stale": True, "never_reported": True} + return snap + + def _globals( + self, + snap: dict[str, dict[str, Any]] | None = None, + max_age_s: float | None = None, + ) -> dict[str, GlobalPoint]: + """World-frame (lat, lon, rel_alt) per robot, skipping stale/missing ones. + + ``max_age_s`` tightens the freshness bound beyond the reporting default; + the guardrail passes GUARDRAIL_MAX_AGE_SEC so it never reasons about + collisions using positions that have had time to become wrong. + """ + snap = snap if snap is not None else self._snapshot() + out: dict[str, GlobalPoint] = {} + for key, state in snap.items(): + if state.get("stale") or not state.get("connected"): + continue + if max_age_s is not None and float(state.get("age_sec", 0.0)) > max_age_s: + continue + g = state.get("global") + if not g: + continue + out[key] = (float(g["lat"]), float(g["lon"]), float(g["rel_alt_m"])) + return out + + def _local_ned(self, key: str) -> tuple[float, float, float] | None: + with self._lock: + state = self._states.get(key) + if not state: + return None + ned = state.get("ned") + if not ned: + return None + return (float(ned[0]), float(ned[1]), float(ned[2])) + + @skill + def list_drones(self) -> str: + """List every robot that has reported in, with its class and state.""" + snap = self._snapshot() + if not snap: + return ( + "No robots have reported yet. The fleet modules may still be " + "connecting — retry in a few seconds." + ) + n_air = len(self._flying_keys()) + n_ground = len(snap) - n_air + header = f"Fleet of {len(snap)} robot(s): {n_air} aircraft" + if n_ground: + header += f", {n_ground} ground" + lines = [header + ":"] + for key in sorted(snap): + st = snap[key] + cls = st.get("robot_class", "multirotor") + parts = [f"class={cls}", f"connected={st.get('connected', False)}"] + # sys_id and armed are aircraft concepts; showing "?" for a ground + # robot reads like missing telemetry rather than "not applicable". + if st.get("sys_id") is not None: + parts.append(f"sys_id={st['sys_id']}") + if st.get("armed") is not None: + parts.append(f"armed={st['armed']}") + parts.append(f"stale={st.get('stale', False)}") + lines.append(f" {key} — " + ", ".join(parts)) + return "\n".join(lines) + + @skill + def fleet_state(self) -> str: + """Report position, armed state, battery %, and pairwise distances for the fleet. + + Call this before commanding motion to confirm the fleet is where you + think it is. ``pairwise_m`` is world-frame 3-D distance (Haversine plus + altitude difference), so the values are real inter-drone separations + regardless of each drone's own NED origin. + """ + snap = self._snapshot() + distances = pairwise_distances(self._globals(snap)) + return json.dumps( + { + "drones": snap, + "pairwise_m": [ + {"a": a, "b": b, "distance": round(d, 2)} for a, b, d in distances + ], + "min_separation_m": self.config.min_separation_m, + }, + indent=2, + ) + + @skill + def count_within(self, drone: str, radius_m: float = 100.0) -> str: + """Count how many other drones are within ``radius_m`` of ``drone``. + + Uses real world-frame 3-D distance (Haversine plus altitude difference). + + Args: + drone: Reference drone key (e.g. ``drone1``). + radius_m: Search radius in meters. + """ + globals_ = self._globals() + if drone not in globals_: + return ( + f"{drone}: no usable global position yet " + f"(known drones: {', '.join(sorted(globals_)) or 'none'})" + ) + ref = globals_[drone] + within: list[str] = [] + for key, point in globals_.items(): + if key == drone: + continue + d = distance_3d_m(ref, point) + if d <= radius_m: + within.append(f"{key} ({d:.1f}m)") + return ( + f"{len(within)} robot(s) within {radius_m:.0f}m of {drone}: " + + (", ".join(sorted(within)) if within else "(none)") + ) + + @skill + def preflight_check(self) -> str: + """Check whether the fleet is safe to fly, and say exactly what is wrong. + + This is the PDF's "hardware safety checklist" as a callable gate. Run it + before every flight, and especially before the first outdoor multi-drone + test. It is read-only — it never commands anything. + + Checks, per drone: + * reporting at all, and not stale + * MAVLink connected + * a global position fix (no fix -> no position commands, no RTL) + * battery at or above ``min_battery_pct`` + * not already armed (an armed drone before takeoff is a surprise) + + And across the fleet: + * every expected drone present + * unique MAV_SYS_IDs — colliding ids make drones unaddressable, and it + is the single most common real-hardware bring-up mistake + * current pairwise separation at or above ``min_separation_m`` + + Returns a PASS/FAIL report. FAIL means do not fly. + """ + snap = self._snapshot() + problems: list[str] = [] + notes: list[str] = [] + + if not snap: + return "FAIL: no robots have reported at all. Is the sim or the telemetry link up?" + + sys_ids: dict[int, list[str]] = {} + for key in sorted(snap): + s = snap[key] + if s.get("never_reported"): + problems.append(f"{key}: expected but never reported") + continue + if s.get("stale"): + problems.append(f"{key}: telemetry stale ({s.get('age_sec')}s old)") + continue + if not s.get("connected"): + problems.append(f"{key}: MAVLink not connected") + continue + if not s.get("global"): + problems.append(f"{key}: no global position fix (GPS) — RTL and goto unavailable") + battery = s.get("battery_pct") + if battery is None: + notes.append(f"{key}: no battery telemetry (placeholder in SITL)") + elif battery < self.config.min_battery_pct: + problems.append( + f"{key}: battery {battery:.0f}% below floor {self.config.min_battery_pct:.0f}%" + ) + if s.get("armed"): + problems.append(f"{key}: already ARMED before preflight") + sid = s.get("sys_id") + if sid is not None: + sys_ids.setdefault(int(sid), []).append(key) + + for sid, keys in sorted(sys_ids.items()): + if len(keys) > 1: + problems.append( + f"MAV_SYS_ID {sid} shared by {', '.join(sorted(keys))} — " + f"set a unique MAV_SYS_ID on each airframe" + ) + + # The separation floor exists for things that fly -- rotor wash and + # in-flight position uncertainty. Two parked ground robots standing a + # metre apart is normal, and failing preflight over it would make this + # check useless on any mixed fleet. + air = set(self._flying_keys()) + air_positions = {k: v for k, v in self._globals(snap).items() if k in air} + for a, b, dist in pairwise_distances(air_positions): + if dist < self.config.min_separation_m: + problems.append( + f"{a} and {b} are {dist:.1f}m apart, below the " + f"{self.config.min_separation_m:.1f}m minimum" + ) + ground_pairs = [ + (a, b, dist) + for a, b, dist in pairwise_distances(self._globals(snap)) + if (a not in air or b not in air) and dist < self.config.min_separation_m + ] + for a, b, dist in ground_pairs: + notes.append(f"{a} and {b} are {dist:.1f}m apart (ground robot, not a flight limit)") + + header = f"Preflight: {len(snap)} robot(s), separation floor " + header += f"{self.config.min_separation_m:.1f}m, battery floor " + header += f"{self.config.min_battery_pct:.0f}%, operating radius " + header += f"{self.config.max_range_m:.0f}m" + lines = [header] + if notes: + lines += ["", "NOTES:"] + [f" - {n}" for n in notes] + if problems: + lines += ["", "FAIL — do not fly:"] + [f" - {p}" for p in problems] + else: + lines += ["", "PASS — all checks clear."] + lines += [ + "", + "Not checkable from here — confirm by hand before a real flight:", + " - RC transmitter bound, and you can flip out of OFFBOARD instantly", + " - PX4 failsafes set: RC loss, datalink loss, low battery, geofence", + " - Geofence configured (GF_ACTION, GF_MAX_HOR_DIST)", + " - Props secured, area clear, spotter briefed", + ] + return "\n".join(lines) + + # ------------------------------------------------------------------ + # Guardrails + # ------------------------------------------------------------------ + + def _check_altitude_cap(self, altitude: float) -> str | None: + cap = self.config.max_altitude_m + if cap is None: + return None + if altitude > cap: + return ( + f"REJECTED: altitude {altitude:.1f}m exceeds fleet cap {cap:.1f}m. " + f"Stay below the cap or raise max_altitude_m." + ) + return None + + def _violates_separation( + self, issuing_key: str, target_n: float, target_e: float, target_d: float + ) -> str | None: + """Predict where ``issuing_key`` would land and check world-frame spacing. + + The target lives in the *issuing drone's* local NED frame, so it is + converted to lat/lon using that drone's current global position plus the + NED delta from its current local position. The prediction is then + Haversine-compared against every other drone's current global position. + + Returns ``"{key} @ {dist}m"`` if the target lands within + ``min_separation_m`` of another drone, otherwise None. + """ + globals_ = self._globals(max_age_s=GUARDRAIL_MAX_AGE_SEC) + issuing_global = globals_.get(issuing_key) + issuing_local = self._local_ned(issuing_key) + # Fail closed. Without a *fresh* position for the issuing vehicle this + # check cannot be performed, and "cannot verify" is not "safe" -- the + # previous behaviour let the command through and hoped PX4 would nack + # it. Telemetry arrives within a second, so refusing costs a retry. + if issuing_global is None or issuing_local is None: + return ( + f"no position for {issuing_key} fresher than " + f"{GUARDRAIL_MAX_AGE_SEC:.2f}s — cannot verify separation" + ) + dn = target_n - issuing_local[0] + de = target_e - issuing_local[1] + pred_lat, pred_lon = offset_latlon(issuing_global[0], issuing_global[1], dn, de) + pred_world = (pred_lat, pred_lon, -target_d) # NED down -> altitude up + for key, point in globals_.items(): + if key == issuing_key: + continue + dist = distance_3d_m(pred_world, point) + if dist < self.config.min_separation_m: + return f"{key} @ {dist:.2f}m" + return None + + # ------------------------------------------------------------------ + # Command bus + # ------------------------------------------------------------------ + + def _send(self, target: str, action: str, repeats: int = 1, **args: Any) -> None: + """Broadcast one command on the swarm bus. + + ``repeats`` exists for the safety verbs: the bus is best-effort, and a + single lost broadcast once left one of two drones hovering armed after + ``rtl_all`` while its twin landed (the module's handler never saw the + message; an individual retry worked instantly). rtl/land/hold/ + emergency/kill are all idempotent, so repeating them is free insurance + against exactly that. Non-idempotent commands keep repeats=1. + """ + for i in range(max(1, repeats)): + with self._lock: + self._seq += 1 + seq = self._seq + payload = {"seq": seq, "target": target, "action": action, "args": args} + self.swarm_cmd.publish(String(json.dumps(payload))) + if i + 1 < repeats: + time.sleep(0.15) + + # What each fleet safety verb must produce in a drone's telemetry. A drone + # matching none of these a few seconds after the broadcast did not hear it. + _VERB_OUTCOME = { + "rtl": lambda st: st.get("mode") in ("AUTO.RTL", "AUTO.LAND") or st.get("armed") is False, + "land": lambda st: st.get("mode") == "AUTO.LAND" or st.get("armed") is False, + "hold": lambda st: st.get("mode") == "AUTO.LOITER" or st.get("armed") is False, + "emergency_land": lambda st: st.get("mode") == "AUTO.LAND" or st.get("armed") is False, + "kill": lambda st: st.get("armed") is False, + } + # Wall-clock settle time before checking; telemetry streams at wall rate + # regardless of the sim's speed, so this is the correct clock. + _verify_delay_s = 4.0 + _verify_rounds = 2 + + def _verify_fleet_command(self, action: str, keys: list[str]) -> None: + """Background delivery check for a broadcast safety verb. + + The bus is best-effort: a lost rtl_all once left one of two drones + hovering armed while its twin landed, and the coordinator had already + reported "dispatched". Dispatched is not delivered -- so after each + safety broadcast this waits, reads every aircraft's reported mode, and + RE-SENDS, individually addressed, to any that did not comply. Stops + after _verify_rounds; a drone still non-compliant then is logged as + needing the operator. + """ + outcome = self._VERB_OUTCOME.get(action) + if outcome is None: + return + pending = list(keys) + for round_no in range(1, self._verify_rounds + 1): + time.sleep(self._verify_delay_s) + snap = self._snapshot() + still = [] + for key in pending: + st = snap.get(key) or {} + # kill verifies on `armed` alone; every other verb needs the + # mode field. Without it we cannot distinguish "did not hear" + # from "mid-descent", so we say we cannot verify rather than + # spam re-sends at a healthy drone. + needs_mode = action != "kill" + if (needs_mode and st.get("mode") is None) or ( + not needs_mode and st.get("armed") is None + ): + logger.warning( + f"[verify:{action}] {key}: telemetry lacks the field needed to " + "verify delivery -- cannot confirm" + ) + continue + if not outcome(st): + still.append(key) + if not still: + if round_no > 1: + logger.info(f"[verify:{action}] all stragglers complied after re-send") + return + for key in still: + logger.warning( + f"[verify:{action}] {key} did not comply (mode=" + f"{(snap.get(key) or {}).get('mode')}) -- re-sending individually" + ) + self._send(key, action, repeats=2) + pending = still + logger.error( + f"[verify:{action}] STILL non-compliant after {self._verify_rounds} rounds: " + f"{', '.join(pending)} -- operator action needed (try the per-drone skill)" + ) + + def _start_fleet_verify(self, action: str, keys: list[str]) -> None: + threading.Thread( + target=self._verify_fleet_command, + args=(action, keys), + name=f"verify-{action}", + daemon=True, + ).start() + + def _known_keys(self, robot_class: str | None = None) -> list[str]: + """Connected robots, optionally restricted to one class. + + Robots that predate the ``robot_class`` field are treated as multirotors, + which is what they were. + """ + snap = self._snapshot() + out = [] + for key, state in snap.items(): + if not state.get("connected"): + continue + if robot_class is not None and state.get("robot_class", "multirotor") != robot_class: + continue + out.append(key) + return sorted(out) + + def _check_range(self, north: float, east: float, what: str) -> str | None: + """Reject a waypoint outside the operating radius, or None if fine.""" + r = math.hypot(north, east) + if r > self.config.max_range_m: + return ( + f"REJECTED {what}: ({north:.0f}, {east:.0f}) is {r:.0f} m out, " + f"beyond the {self.config.max_range_m:.0f} m operating radius. " + "Pick a closer target." + ) + return None + + def _flying_keys(self) -> list[str]: + """Robots that can actually hold an altitude. + + Maneuvers that assign a 3-D lane -- grid_sweep, line_formation, + investigate -- must use this rather than every reporting robot. Handing a + quadruped an air lane silently loses that share of the coverage, and in + `investigate` it consumes one of the requested units. + """ + return self._known_keys(robot_class="multirotor") + + def _airborne_keys(self) -> list[str]: + """Aircraft that are actually armed and flying. + + Position maneuvers engage OFFBOARD. Doing that to a disarmed vehicle + sitting on the ground is not just useless: PX4 trips an OFFBOARD + failsafe and then refuses the *next* arm request with "Resolve system + health failures first", so a stray grid_sweep before takeoff silently + breaks takeoff afterwards. + """ + snap = self._snapshot() + return sorted( + k + for k in self._flying_keys() + if snap.get(k, {}).get("armed") + ) + + def _needs_airborne(self, maneuver: str) -> str | None: + """Refuse an air maneuver when nothing is flying, and say what to do.""" + if self._airborne_keys(): + return None + if not self._flying_keys(): + return "No aircraft have reported in yet." + return ( + f"REJECTED: no aircraft are armed, so {maneuver} would engage OFFBOARD on " + f"grounded vehicles and trip a failsafe that blocks the next arm. " + f"Call takeoff_all first." + ) + + def _ground_note(self) -> str: + """Trailing note naming ground robots skipped by an air maneuver.""" + ground = [k for k in self._known_keys() if k not in set(self._flying_keys())] + return f"\n (ground robots not assigned an air lane: {', '.join(ground)})" if ground else "" + + # ------------------------------------------------------------------ + # Aggregate skills + # ------------------------------------------------------------------ + + @skill + def takeoff_all(self, altitude: float = 3.0) -> str: + """Arm and take off every drone to ``altitude`` (m above home) simultaneously.""" + cap_err = self._check_altitude_cap(altitude) + if cap_err: + return cap_err + # Dispatch stays a broadcast: a drone whose telemetry is momentarily + # stale still hears it, where a filtered list would silently skip it. + # Ground robots ignore air actions. The REPORT, though, must name only + # aircraft -- listing a quadruped under "takeoff" reads as though it + # left the ground, which is the kind of false success that wastes an + # hour of debugging. + keys = self._flying_keys() + if not keys: + return "REJECTED: no aircraft in the fleet to take off." + self._ground_note() + self._send("all", "takeoff", altitude=altitude) + return ( + f"takeoff_all to {altitude}m dispatched to {len(keys)} aircraft: " + f"{', '.join(keys)}" + self._ground_note() + ) + + @skill + def land_all(self) -> str: + """Land every drone where it is, simultaneously.""" + keys = self._flying_keys() + if not keys: + return "REJECTED: no aircraft in the fleet to land." + self._ground_note() + self._send("all", "land", repeats=3) + self._start_fleet_verify("land", self._flying_keys()) + return ( + f"land_all dispatched to {len(keys)} aircraft: {', '.join(keys)}" + + self._ground_note() + ) + + @skill + def rtl_all(self) -> str: + """Return every drone to launch and land — the PDF's "Return to launch".""" + keys = self._flying_keys() + if not keys: + return "REJECTED: no aircraft in the fleet to return." + self._ground_note() + self._send("all", "rtl", repeats=3) + self._start_fleet_verify("rtl", self._flying_keys()) + return ( + f"rtl_all dispatched to {len(keys)} aircraft: {', '.join(keys)}" + + self._ground_note() + ) + + # ------------------------------------------------------------------ + # Staged missions + # ------------------------------------------------------------------ + + def _legged_keys(self) -> list[str]: + return self._known_keys(robot_class="legged") + + def _dog_at(self, key: str, north: float, east: float, radius_m: float = 1.5) -> bool: + st = self._snapshot().get(key) or {} + ned = st.get("ned") + if not ned: + return False + return math.hypot(north - ned[0], east - ned[1]) <= radius_m + + def _set_mission(self, stage: str, detail: str) -> None: + self.__dict__["_mission"] = {"stage": stage, "detail": detail, "ts": time.time()} + logger.info(f"[mission] {stage}: {detail}") + + @skill + def sweep_then_ground( + self, + corner_b_north: float = 40.0, + corner_b_east: float = 30.0, + altitude: float = 8.0, + ground_delay_s: float = 30.0, + dog: str = "", + ) -> str: + """Air sweep first, ground robot in after a delay -- the staged demo. + + The drones fly a boustrophedon over (0,0)..(corner_b) immediately. After + ``ground_delay_s`` (wall-clock seconds -- at ~10x realtime that is far + longer in sim time, deliberately, since the operator watches the wall + clock) the ground robot walks a straight transect up the middle of the + swept area -- entry, centre, far edge -- then returns to where it + started. Progress is readable at any time via ``mission_status``, and + ``emergency_land_all`` / ``kill_all`` abort the ground stage. + + The default 40 x 30 m rectangle matches the marked field in the sim + world, so the whole mission is legible in the viewer. + """ + prior = self.__dict__.get("_mission_thread") + if prior is not None and prior.is_alive(): + return "REJECTED: a staged mission is already running. mission_status to watch it." + # Reuse grid_sweep wholesale: its altitude, range, and airborne gates + # are the authority, and its rejection strings say what to fix. + sweep = self.grid_sweep( + corner_b_north=corner_b_north, corner_b_east=corner_b_east, altitude=altitude + ) + if sweep.startswith("REJECTED"): + return sweep + dogs = self._legged_keys() + if dog: + if dog not in dogs: + return f"Unknown ground robot '{dog}'. Known: {', '.join(dogs) or 'none'}" + picked = dog + elif dogs: + picked = dogs[0] + else: + return ( + "REJECTED: no ground robot in the fleet -- sweep dispatched to nobody. " + "Use grid_sweep for an air-only sweep." + ) + start = (self._snapshot().get(picked) or {}).get("ned") or [0.0, 0.0, 0.0] + mid_e = corner_b_east / 2.0 + route = [ + (0.0, mid_e, "entry"), + (corner_b_north / 2.0, mid_e, "centre"), + (corner_b_north, mid_e, "far edge"), + (start[0], start[1], "home"), + ] + abort = threading.Event() + self.__dict__["_mission_abort"] = abort + timeout_s = float(self.__dict__.get("_ground_wait_timeout_s", 180.0)) + + def _run() -> None: + self._set_mission( + "air sweep", f"drones sweeping; ground in {ground_delay_s:.0f}s" + ) + if abort.wait(timeout=max(0.0, ground_delay_s)): + self._set_mission("aborted", "before ground stage") + return + for wn, we, label in route: + if abort.is_set(): + self._set_mission("aborted", f"during ground stage at {label}") + return + self._set_mission("ground transect", f"{picked} -> {label} ({wn:.0f}, {we:.0f})") + self._send(picked, "ground_goto", north=wn, east=we) + deadline = time.time() + timeout_s + while time.time() < deadline and not abort.is_set(): + if self._dog_at(picked, wn, we): + break + if (self._snapshot().get(picked) or {}).get("fallen"): + self._set_mission("failed", f"{picked} fell en route to {label}") + return + time.sleep(0.5) + else: + if not abort.is_set(): + logger.warning(f"[mission] {picked} timed out reaching {label}; continuing") + self._set_mission("complete", f"{picked} back at home; drones hold at lane ends") + + t = threading.Thread(target=_run, daemon=True, name="staged-mission") + self.__dict__["_mission_thread"] = t + t.start() + return ( + f"Staged sweep started over (0,0)..({corner_b_north:.0f},{corner_b_east:.0f}):\n" + f" now : {len(self._flying_keys())} aircraft sweeping at {altitude:.0f} m\n" + f" +{ground_delay_s:.0f}s : {picked} walks the centre transect " + f"(entry -> centre -> far edge -> home)\n" + " drones hold at their lane ends when done -- rtl_all to bring them back.\n" + " mission_status shows the current stage." + ) + + @skill + def mission_status(self) -> str: + """Current stage of the staged mission, if any. Read-only.""" + m = self.__dict__.get("_mission") + if not m: + return "No staged mission has run." + age = time.time() - m["ts"] + return f"[{m['stage']}] {m['detail']} (stage entered {age:.0f}s ago)" + + @skill + def boundaries(self) -> str: + """Report every active control boundary and which layer enforces it. + + Read-only. Two layers exist and they catch different things: + dispatch-time checks reject a bad *command* before anything moves; + the PX4 geofence catches what dispatch cannot see -- velocity-command + drift, wind in a future HITL setup, or a maneuver mid-flight. + """ + cap = ( + f"{self.config.max_altitude_m:.0f} m" + if self.config.max_altitude_m is not None + else "none set (each drone still enforces its own)" + ) + return ( + "Fleet control boundaries:\n" + f" operating radius : {self.config.max_range_m:.0f} m from home " + "(dispatch-time: goto_drone, investigate, grid_sweep, line_formation, " + "and each dog's goto)\n" + f" altitude ceiling : {cap} (dispatch-time)\n" + f" separation floor : {self.config.min_separation_m:.1f} m between aircraft " + "(dispatch-time admission; NOT in-flight collision avoidance)\n" + " PX4 geofence : GF_ACTION=hold at GF_MAX_HOR_DIST/GF_MAX_VER_DIST " + "(in-flight backstop for AIRCRAFT, written by sim_params.py -- " + "catches set_velocity drift that dispatch checks cannot)\n" + " ground fence : each legged robot halts itself when a walk " + "carries it past the operating radius (in-motion backstop for GROUND " + "robots; walking back in is allowed, further escape re-trips it)\n" + " ground proximity : two ground robots converging within 0.6 m are " + "both halted by the simulator (in-motion; re-arms at 1.2 m). Catches " + "what no dispatch check can: robots ALREADY moving toward each other\n" + " GCS heartbeat : PX4 refuses to arm without DimOS alive, and " + "returns+lands every aircraft on its own if DimOS dies mid-flight " + "(NAV_DLL_ACT=2)\n" + " airborne gate : position maneuvers refuse until takeoff_all " + "(a grounded OFFBOARD trips a failsafe that blocks the next arm)" + ) + + @skill + def hold_all(self) -> str: + """Put every drone into AUTO.LOITER (hover in place).""" + self._send("all", "hold", repeats=3) + self._start_fleet_verify("hold", self._flying_keys()) + keys = self._flying_keys() + return ( + f"hold_all dispatched to {len(keys)} aircraft: {', '.join(keys)}" + + self._ground_note() + ) + + @skill + def emergency_land_all(self) -> str: + """Land every drone immediately, regardless of state. + + Safer than ``kill_all`` — each drone descends under its own control + instead of falling. Use when something is wrong but the autopilots are + still responsive. ``kill_all`` is the next escalation step. + """ + ab = self.__dict__.get("_mission_abort") + if ab is not None: + ab.set() + # Broadcast on purpose -- an emergency must reach anything listening, + # including a drone whose class or telemetry is momentarily missing. + self._send("all", "emergency_land", repeats=3) + self._start_fleet_verify("emergency_land", self._flying_keys()) + return f"EMERGENCY LAND broadcast; {len(self._flying_keys())} aircraft known" + + @skill + def kill_all(self) -> str: + """Force-disarm every drone — the big red button. They fall. + + Use ONLY when crashing is preferable to what the fleet is about to do. + Try ``emergency_land_all`` first if the autopilots still respond. + """ + ab = self.__dict__.get("_mission_abort") + if ab is not None: + ab.set() + # Broadcast on purpose: see emergency_land_all. Never narrow the big + # red button to a filtered list. + self._send("all", "kill", repeats=3) + self._start_fleet_verify("kill", self._flying_keys()) + logger.warning("KILL_ALL broadcast on swarm_cmd") + return f"KILL_ALL broadcast; {len(self._flying_keys())} aircraft known. Motors stopped." + + # ------------------------------------------------------------------ + # Demo maneuvers + # ------------------------------------------------------------------ + + @skill + def line_formation( + self, + center_north: float = 0.0, + center_east: float = 0.0, + altitude: float = 5.0, + spacing_m: float = 4.0, + heading_deg: float = 90.0, + ) -> str: + """Position all drones on a line at fixed altitude and spacing. + + The PDF's "Return to line formation." The line is centered at + ``(center_north, center_east)`` and oriented along ``heading_deg`` (NED + bearing, 0 = north, 90 = east). Spacing is forced to at least + ``min_separation_m + 0.5`` so the formation cannot violate the guardrail. + """ + cap_err = self._check_altitude_cap(altitude) + if cap_err: + return cap_err + blocked = self._needs_airborne("line_formation") + if blocked: + return blocked + keys = self._airborne_keys() + spacing = max(spacing_m, self.config.min_separation_m + 0.5) + bearing = math.radians(heading_deg) + n = len(keys) + half = (n - 1) / 2.0 * spacing + for sign in (1.0, -1.0): + end_n = center_north + sign * half * math.cos(bearing) + end_e = center_east + sign * half * math.sin(bearing) + range_err = self._check_range(end_n, end_e, "line_formation (line end)") + if range_err: + return range_err + results: list[str] = [] + for i, key in enumerate(keys): + offset = (i - (n - 1) / 2.0) * spacing + tn = center_north + offset * math.cos(bearing) + te = center_east + offset * math.sin(bearing) + self._send(key, "goto_ned", north=tn, east=te, down=-altitude, yaw_rad=bearing) + results.append(f"{key}: → ({tn:.1f}, {te:.1f}, alt={altitude:.1f}m)") + return ( + f"Line formation (spacing {spacing:.1f}m):\n " + + "\n ".join(results) + + self._ground_note() + ) + + @skill + def grid_sweep( + self, + corner_a_north: float = 0.0, + corner_a_east: float = 0.0, + corner_b_north: float = 30.0, + corner_b_east: float = 30.0, + altitude: float = 8.0, + sub_lane_spacing_m: float = 5.0, + ) -> str: + """Run a boustrophedon ("lawnmower") sweep with the fleet in parallel strips. + + The PDF's "Take your team and sweep this grid." The rectangle A→B is + split into one equal-width east-axis strip per drone. Each drone walks a + zigzag inside its own strip, so no two drones ever share a lane. + + Args: + corner_a_north, corner_a_east: one corner of the sweep area, in each + drone's local NED frame (meters relative to its own home). + corner_b_north, corner_b_east: the opposite corner. + altitude: sweep altitude (m, positive up). + sub_lane_spacing_m: distance between zigzag passes inside a strip. + Smaller = denser coverage, longer flight. + + Guardrails: strips are separated by a ``min_separation_m / 2`` buffer at + each edge, so neighboring drones stay at least ``min_separation_m`` apart + even at their closest passes. Rejected outright if the strips would be + too narrow for that buffer. + + The sweep does NOT auto-return; drones hover at their lane ends. Call + ``rtl_all`` afterwards if the intent was "search then come back". + """ + cap_err = self._check_altitude_cap(altitude) + if cap_err: + return cap_err + blocked = self._needs_airborne("grid_sweep") + if blocked: + return blocked + keys = self._airborne_keys() + + north_min = min(corner_a_north, corner_b_north) + north_max = max(corner_a_north, corner_b_north) + east_min = min(corner_a_east, corner_b_east) + east_max = max(corner_a_east, corner_b_east) + # The farthest point of the rectangle must stay inside the fence. + far_n = max(abs(north_min), abs(north_max)) + far_e = max(abs(east_min), abs(east_max)) + range_err = self._check_range(far_n, far_e, "grid_sweep (far corner)") + if range_err: + return range_err + n = len(keys) + strip_width = (east_max - east_min) / n + + edge_buffer = self.config.min_separation_m / 2.0 + usable_strip = strip_width - 2 * edge_buffer + if usable_strip < 0: + return ( + f"REJECTED: strip width {strip_width:.1f}m too narrow for a " + f"{self.config.min_separation_m}m separation buffer across {n} drones. " + f"Widen the area or use fewer drones." + ) + + results: list[str] = [] + for i, key in enumerate(keys): + strip_e_lo = east_min + i * strip_width + edge_buffer + strip_e_hi = east_min + (i + 1) * strip_width - edge_buffer + + sub_e: list[float] = [] + if usable_strip < 0.5 or sub_lane_spacing_m <= 0: + sub_e = [(strip_e_lo + strip_e_hi) / 2.0] + else: + e = strip_e_lo + while e <= strip_e_hi + 1e-3: + sub_e.append(e) + e += sub_lane_spacing_m + if sub_e[-1] < strip_e_hi - 1e-3: + sub_e.append(strip_e_hi) + + # Even sub-lanes walk south→north, odd walk north→south: a + # continuous zigzag with no teleports between passes. + d = -altitude + path: list[list[float]] = [] + for j, e in enumerate(sub_e): + if j % 2 == 0: + path.append([north_min, e, d, 0.0]) + path.append([north_max, e, d, 0.0]) + else: + path.append([north_max, e, d, 0.0]) + path.append([north_min, e, d, 0.0]) + + self._send(key, "path", path=path, arrival_radius_m=2.0) + results.append( + f"{key}: strip east=[{strip_e_lo:.1f}, {strip_e_hi:.1f}]m, " + f"{len(sub_e)} pass(es), {len(path)} waypoints, alt={altitude:.1f}m" + ) + return ( + "Grid sweep dispatched (boustrophedon):\n " + + "\n ".join(results) + + self._ground_note() + ) + + @skill + def investigate( + self, + north: float = 0.0, + east: float = 0.0, + altitude: float = 5.0, + num_drones: int = 2, + drones: str = "", + ) -> str: + """Task N drones to converge on one coordinate (local NED). + + The PDF's "Task two units to investigate this coordinate." The chosen + drones are staggered north-south around the target at + 1.5x ``min_separation_m`` intervals so they do not collide, and each + assignment is still checked against the spacing guardrail. + + Args: + north: Target north position (m, local NED relative to home). + east: Target east position (m). + altitude: Target altitude (m, positive up). + num_drones: How many drones to send. Ignored when ``drones`` is set. + drones: CSV of drone keys (e.g. ``"drone1,drone3"``) to send instead. + """ + cap_err = self._check_altitude_cap(altitude) + if cap_err: + return cap_err + blocked = self._needs_airborne("investigate") + if blocked: + return blocked + known = self._airborne_keys() + + if drones.strip(): + picked = [d.strip() for d in drones.split(",") if d.strip()] + unknown = [d for d in picked if d not in known] + if unknown: + return f"Unknown drone(s): {', '.join(unknown)}. Known: {', '.join(known)}" + else: + picked = known[: max(1, min(num_drones, len(known)))] + + spacing = self.config.min_separation_m * 1.5 + k = len(picked) + # The stagger pushes the outermost drone (k-1)/2 slots past the target; + # that farthest assignment is the one the fence has to admit. + far_n = abs(north) + (k - 1) / 2.0 * spacing + range_err = self._check_range(math.copysign(far_n, north or 1.0), east, "investigate") + if range_err: + return range_err + results: list[str] = [] + for i, key in enumerate(picked): + offset = (i - (k - 1) / 2.0) * spacing + tn = north + offset + te = east + d = -altitude + violation = self._violates_separation(key, tn, te, d) + if violation is not None: + results.append( + f"{key}: REJECTED ({tn:.1f},{te:.1f},alt={altitude:.1f}) — " + f"within {self.config.min_separation_m}m of {violation}" + ) + continue + self._send(key, "goto_ned", north=tn, east=te, down=d, yaw_rad=0.0) + results.append(f"{key}: → ({tn:.1f}, {te:.1f}, alt={altitude:.1f}m)") + return "Investigate dispatched:\n " + "\n ".join(results) + self._ground_note() + + @skill + def goto_drone( + self, + drone: str, + north: float = 0.0, + east: float = 0.0, + altitude: float = 3.0, + yaw_deg: float = 0.0, + ) -> str: + """Send one drone to a local-NED waypoint, with the spacing guardrail applied. + + Prefer this over the drone's own ``goto`` when other drones are + airborne: only the coordinator can see the whole fleet, so only it can + reject a target that would land too close to someone else. + """ + cap_err = self._check_altitude_cap(altitude) + if cap_err: + return cap_err + range_err = self._check_range(north, east, "goto_drone") + if range_err: + return range_err + blocked = self._needs_airborne("goto_drone") + if blocked: + return blocked + known = self._airborne_keys() + if drone not in known: + if drone in self._known_keys(): + return ( + f"{drone} is a ground robot; goto_drone assigns an altitude. " + f"Use its own namespaced skills instead." + ) + return f"Unknown drone: {drone!r}. Known aircraft: {', '.join(known) or 'none'}" + d = -altitude + violation = self._violates_separation(drone, north, east, d) + if violation is not None: + return ( + f"{drone}: REJECTED goto({north:.1f}, {east:.1f}, alt={altitude:.1f}) — " + f"target within {self.config.min_separation_m}m of {violation}" + ) + self._send(drone, "goto_ned", north=north, east=east, down=d, yaw_rad=math.radians(yaw_deg)) + return f"{drone}: goto NED ({north:.1f}, {east:.1f}, alt={altitude:.1f}m), yaw={yaw_deg:.0f}°" + + # ------------------------------------------------------------------ + # Map-aware skills + # ------------------------------------------------------------------ + + def _get_gmaps_client(self) -> Any: + """Lazily build and cache a Google Maps client, or None if unavailable.""" + if self._gmaps_client is False: + return None + if self._gmaps_client is not None: + return self._gmaps_client + api_key = os.getenv("GOOGLE_MAPS_API_KEY") + if not api_key: + logger.info("GOOGLE_MAPS_API_KEY not set — map skills disabled") + self._gmaps_client = False + return None + try: + import googlemaps # type: ignore[import-untyped] + + self._gmaps_client = googlemaps.Client(key=api_key) + except Exception as e: + logger.info(f"googlemaps unavailable — map skills disabled: {e}") + self._gmaps_client = False + return None + return self._gmaps_client + + @skill + def find_place_near(self, drone: str = "", query: str = "") -> str: + """Find the nearest place matching ``query`` to a drone's current position. + + The PDF's "go to the nearest lake" step 1. This only *resolves* a place + to coordinates — it never dispatches. Present the result to the user and + get explicit confirmation before calling ``goto_drone_global``. + + Args: + drone: Reference drone key. Defaults to any connected drone. + query: Free-text place query, e.g. ``"lake"`` or ``"coastline"``. + """ + client = self._get_gmaps_client() + if client is None: + return ( + "Google Maps unavailable — the daemon needs GOOGLE_MAPS_API_KEY " + "in its environment and the `googlemaps` package installed." + ) + if not query.strip(): + return "Provide a query, e.g. find_place_near(drone='drone1', query='lake')." + globals_ = self._globals() + if not globals_: + return "No drone has a usable global position yet." + key = drone if drone in globals_ else next(iter(sorted(globals_))) + lat, lon, _alt = globals_[key] + try: + res = client.places_nearby(location=(lat, lon), keyword=query, rank_by="distance") + except Exception as e: + return f"Google Maps lookup failed: {e}" + results = res.get("results") or [] + if not results: + return f"No place matching {query!r} found near {key}." + top = results[0] + loc = top["geometry"]["location"] + dist = distance_3d_m((lat, lon, 0.0), (loc["lat"], loc["lng"], 0.0)) + return json.dumps( + { + "reference_drone": key, + "name": top.get("name"), + "address": top.get("vicinity"), + "lat": loc["lat"], + "lon": loc["lng"], + "distance_m": round(dist, 1), + "next_step": ( + "Ask the user to confirm before calling goto_drone_global with " + "these coordinates." + ), + }, + indent=2, + ) + + @skill + def goto_drone_global( + self, + drone: str, + lat: float, + lon: float, + altitude: float = 20.0, + ) -> str: + """Fly one drone to a global lat/lon at ``altitude`` (m above its home). + + Only call this AFTER the user has explicitly confirmed the destination — + map-derived goals must never dispatch unconfirmed. + """ + cap_err = self._check_altitude_cap(altitude) + if cap_err: + return cap_err + globals_ = self._globals() + if drone not in globals_: + return f"{drone}: no usable global position yet — cannot compute the NED offset." + cur_lat, cur_lon, _ = globals_[drone] + local = self._local_ned(drone) + if local is None: + return f"{drone}: no local NED yet." + # Convert the global target into this drone's local NED frame. + from dimos.robot.drone.px4_geo import WGS84_EQUATORIAL_M + + dn = math.radians(lat - cur_lat) * WGS84_EQUATORIAL_M + de = math.radians(lon - cur_lon) * WGS84_EQUATORIAL_M * math.cos(math.radians(cur_lat)) + target_n = local[0] + dn + target_e = local[1] + de + violation = self._violates_separation(drone, target_n, target_e, -altitude) + if violation is not None: + return ( + f"{drone}: REJECTED global goto — target within " + f"{self.config.min_separation_m}m of {violation}" + ) + self._send(drone, "goto_ned", north=target_n, east=target_e, down=-altitude, yaw_rad=0.0) + return ( + f"{drone}: goto global ({lat:.6f}, {lon:.6f}) at {altitude:.1f}m " + f"→ local NED ({target_n:.1f}, {target_e:.1f})" + ) + + +__all__ = ["STATE_STALE_AFTER_SEC", "SwarmCoordinator", "SwarmCoordinatorConfig"] diff --git a/dimos/robot/drone/test_px4_swarm.py b/dimos/robot/drone/test_px4_swarm.py new file mode 100644 index 0000000000..ee32bcb0b0 --- /dev/null +++ b/dimos/robot/drone/test_px4_swarm.py @@ -0,0 +1,1145 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the namespaced PX4 swarm. + +Covers the three things the architecture split is supposed to guarantee: + +1. Blueprint composition — N drones really do land in N namespaces, per-drone + streams get prefixed, and the fleet bus (``drone_state``/``swarm_cmd``) + stays global so it can cross the boundary. +2. The command bus — the coordinator's broadcasts reach the right drones and + only the right drones. +3. The fleet logic that used to live in ``Px4SitlFleetModule`` — spacing + guardrail, altitude caps, sweep geometry, formations — still behaves. + +None of this needs a live MAVLink link or a running sim. +""" + +from __future__ import annotations + +from itertools import pairwise +import json +import math +import threading +import time +from types import SimpleNamespace +from typing import Any + +from dimos_lcm.std_msgs import String +import pytest + +from dimos.robot.drone.px4_drone_module import Px4DroneModule +from dimos.robot.drone.px4_geo import distance_3d_m, offset_latlon, pairwise_distances +from dimos.robot.drone.px4_swarm_coordinator import STATE_STALE_AFTER_SEC, SwarmCoordinator + +# A reference point with easy arithmetic (Zurich Irchel is PX4's SITL default home). +HOME_LAT, HOME_LON = 47.397742, 8.545594 + + +# --------------------------------------------------------------------------- +# 1. Blueprint composition / namespacing +# --------------------------------------------------------------------------- + + +def _swarm_blueprint(n: int = 3): + from dimos.robot.drone.blueprints.basic.drone_px4_sitl_fleet_mcp import px4_sitl_swarm + from dimos.robot.drone.px4_sitl_fleet_config import get_px4_sitl_fleet_configs + + return px4_sitl_swarm(configs=get_px4_sitl_fleet_configs(fleet_size=n)) + + +def test_each_drone_gets_its_own_namespaced_instance(): + bp = _swarm_blueprint(3) + names = {a.name for a in bp.blueprints} + assert names == { + "swarmcoordinator", + "drone1/px4dronemodule", + "drone2/px4dronemodule", + "drone3/px4dronemodule", + } + + +def test_fleet_size_is_dynamic(): + assert len([a for a in _swarm_blueprint(2).blueprints if "px4dronemodule" in a.name]) == 2 + assert len([a for a in _swarm_blueprint(5).blueprints if "px4dronemodule" in a.name]) == 5 + + +def test_per_drone_streams_are_prefixed(): + """cmd_vel and odom must NOT be shared, or one tracker would fly every drone.""" + remap = _swarm_blueprint(3).remapping_map + for i in (1, 2, 3): + assert remap[(f"drone{i}/px4dronemodule", "cmd_vel")] == f"drone{i}/cmd_vel" + assert remap[(f"drone{i}/px4dronemodule", "odom")] == f"drone{i}/odom" + + +def test_fleet_bus_streams_stay_global(): + """drone_state/swarm_cmd are exposed, so they must have no namespace remap.""" + remap = _swarm_blueprint(3).remapping_map + for i in (1, 2, 3): + assert (f"drone{i}/px4dronemodule", "drone_state") not in remap + assert (f"drone{i}/px4dronemodule", "swarm_cmd") not in remap + + +def test_each_drone_gets_its_own_connection_string(): + bp = _swarm_blueprint(3) + conns = { + a.name: a.kwargs["connection_string"] for a in bp.blueprints if "px4dronemodule" in a.name + } + assert conns == { + "drone1/px4dronemodule": "udp:127.0.0.1:14540", + "drone2/px4dronemodule": "udp:127.0.0.1:14541", + "drone3/px4dronemodule": "udp:127.0.0.1:14542", + } + + +def test_config_keys_are_unique_per_instance(): + """Namespacing is what makes N instances of one class configurable separately.""" + from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser + + parsed = BlueprintConfigParser(_swarm_blueprint(3)).parse( + environ={}, overrides={"g": {"viewer": "none"}} + ) + assert { + "drone1/px4dronemodule", + "drone2/px4dronemodule", + "drone3/px4dronemodule", + "swarmcoordinator", + } <= set(parsed.module_configs) + + +def test_a_single_drone_can_be_reconfigured_without_touching_the_others(): + """`-o drone2/px4dronemodule.max_altitude_m=15` must hit only drone2.""" + from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser + + parsed = BlueprintConfigParser(_swarm_blueprint(3)).parse( + environ={}, + overrides={ + "g": {"viewer": "none"}, + "drone2_px4dronemodule": {"max_altitude_m": 15.0}, + }, + ) + assert parsed.module_kwargs("drone2/px4dronemodule")["max_altitude_m"] == 15.0 + assert parsed.module_kwargs("drone1/px4dronemodule").get("max_altitude_m") != 15.0 + + +# --------------------------------------------------------------------------- +# 2. Drone identity + command bus +# --------------------------------------------------------------------------- + + +def _bare_drone(instance_name: str | None = "drone2/px4dronemodule", max_alt: float | None = None): + m = object.__new__(Px4DroneModule) + m.config = SimpleNamespace(instance_name=instance_name, max_altitude_m=max_alt) + m._follow_enabled = False + m._follow_lock_altitude = True + m.connection = None + return m + + +def test_drone_key_comes_from_the_namespace(): + assert _bare_drone("drone2/px4dronemodule").drone_key == "drone2" + assert _bare_drone("drone11/px4dronemodule").drone_key == "drone11" + + +def test_drone_key_falls_back_when_not_namespaced(): + assert _bare_drone(None).drone_key == "drone1" + + +def _dispatched(drone: Px4DroneModule) -> list[tuple[str, dict[str, Any]]]: + """Record what _on_swarm_cmd actually dispatches, without touching MAVLink.""" + seen: list[tuple[str, dict[str, Any]]] = [] + drone._dispatch = lambda action, args: ( # type: ignore[method-assign] + seen.append((action, args)) or "ok" + ) + return seen + + +def _cmd(target: str, action: str, **args: Any) -> String: + return String(json.dumps({"seq": 1, "target": target, "action": action, "args": args})) + + +def test_drone_acts_on_commands_addressed_to_it(): + d = _bare_drone("drone2/px4dronemodule") + seen = _dispatched(d) + d._on_swarm_cmd(_cmd("drone2", "takeoff", altitude=5.0)) + assert seen == [("takeoff", {"altitude": 5.0})] + + +def test_drone_ignores_commands_for_other_drones(): + d = _bare_drone("drone2/px4dronemodule") + seen = _dispatched(d) + d._on_swarm_cmd(_cmd("drone1", "takeoff", altitude=5.0)) + assert seen == [] + + +def test_drone_acts_on_broadcast_commands(): + d = _bare_drone("drone2/px4dronemodule") + seen = _dispatched(d) + d._on_swarm_cmd(_cmd("all", "rtl")) + assert seen == [("rtl", {})] + + +def test_malformed_command_is_ignored_not_raised(): + d = _bare_drone("drone2/px4dronemodule") + _dispatched(d) + d._on_swarm_cmd(String("not json at all")) # must not raise + + +def test_unknown_action_is_reported_not_raised(): + d = _bare_drone("drone2/px4dronemodule") + assert "unknown action" in d._dispatch("fly_to_the_moon", {}) + + +def test_altitude_cap_rejects_and_allows(): + d = _bare_drone("drone1/px4dronemodule", max_alt=30.0) + assert d._check_altitude_cap(10.0) is None + assert "REJECTED" in (d._check_altitude_cap(50.0) or "") + # No cap configured -> nothing is rejected. + assert _bare_drone("drone1/px4dronemodule", max_alt=None)._check_altitude_cap(500.0) is None + + +def test_takeoff_is_blocked_by_the_cap_before_touching_mavlink(): + d = _bare_drone("drone1/px4dronemodule", max_alt=30.0) + d.connection = None # would fail with NOT CONNECTED if the cap did not fire first + assert "REJECTED" in d.takeoff(999.0) + + +# --------------------------------------------------------------------------- +# 3. Coordinator: state aggregation, guardrails, maneuvers +# --------------------------------------------------------------------------- + + +class _FakeOut: + def __init__(self) -> None: + self.published: list[dict[str, Any]] = [] + + def publish(self, msg: String) -> None: + self.published.append(json.loads(msg.data if hasattr(msg, "data") else str(msg))) + + +def _coordinator( + min_separation_m: float = 2.0, + max_altitude_m: float | None = None, + expected: str = "", + min_battery_pct: float = 25.0, + max_range_m: float = 250.0, +) -> SwarmCoordinator: + c = object.__new__(SwarmCoordinator) + c.config = SimpleNamespace( + min_separation_m=min_separation_m, + max_altitude_m=max_altitude_m, + expected_drones=expected, + min_battery_pct=min_battery_pct, + max_range_m=max_range_m, + ) + c._states = {} + c._lock = threading.RLock() + c._seq = 0 + c._gmaps_client = False + c.swarm_cmd = _FakeOut() + + # Delivery verification runs synchronously and instantly in tests: no + # leaked threads, deterministic message counts, and the resend behavior + # itself is covered by its own tests below. + c._verify_delay_s = 0.0 + c._start_fleet_verify = lambda action, keys: c._verify_fleet_command(action, keys) + return c + + +def _report(c: SwarmCoordinator, key: str, n: float, e: float, alt: float, **extra: Any) -> None: + """Feed one drone_state message in, positioned n/e meters from HOME. + + Armed by default: the position maneuvers require an airborne aircraft, so an + unarmed fleet is the exception rather than the norm in these tests. + """ + lat, lon = offset_latlon(HOME_LAT, HOME_LON, n, e) + state = { + "key": key, + "robot_class": "multirotor", + "connected": True, + "armed": True, + "battery_pct": 95.0, + "ned": [n, e, -alt], + "altitude_m": alt, + "global": {"lat": lat, "lon": lon, "rel_alt_m": alt}, + "ts": time.time(), + } + state.update(extra) + c._on_drone_state(String(json.dumps(state))) + + +def test_coordinator_aggregates_state_from_the_bus(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 10, 0, 5) + state = json.loads(c.fleet_state()) + assert set(state["drones"]) == {"drone1", "drone2"} + assert state["drones"]["drone1"]["battery_pct"] == 95.0 + + +def test_pairwise_distances_are_world_frame(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 30, 0, 5) + pairs = json.loads(c.fleet_state())["pairwise_m"] + assert len(pairs) == 1 + assert pairs[0]["distance"] == pytest.approx(30.0, abs=0.5) + + +def test_stale_drones_are_flagged_and_excluded_from_geometry(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 10, 0, 5) + # Backdate drone2 past the staleness horizon. + c._states["drone2"]["ts"] = time.time() - (STATE_STALE_AFTER_SEC + 10) + snap = json.loads(c.fleet_state()) + assert snap["drones"]["drone2"]["stale"] is True + assert snap["pairwise_m"] == [] # a stale drone can't anchor guardrail math + + +def test_expected_but_silent_drones_are_reported(): + c = _coordinator(expected="drone1,drone2,drone3") + _report(c, "drone1", 0, 0, 5) + drones = json.loads(c.fleet_state())["drones"] + assert drones["drone3"]["never_reported"] is True + assert drones["drone3"]["connected"] is False + + +def test_count_within_matches_the_pdf_demo_question(): + """"How many units are within 100 meters of whatever drone".""" + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 50, 0, 5) # inside 100 m + _report(c, "drone3", 500, 0, 5) # outside + out = c.count_within("drone1", 100.0) + assert out.startswith("1 robot(s) within 100m of drone1") + assert "drone2" in out and "drone3" not in out + + +def test_count_within_unknown_drone_is_reported(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + assert "no usable global position" in c.count_within("ghost", 100.0) + + +# --- spacing guardrail ----------------------------------------------------- + + +def test_guardrail_rejects_a_target_too_close_to_another_drone(): + c = _coordinator(min_separation_m=5.0) + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 20, 0, 5) + # Send drone1 to within 1 m of drone2. + out = c.goto_drone("drone1", north=19.0, east=0.0, altitude=5.0) + assert "REJECTED" in out and "drone2" in out + assert c.swarm_cmd.published == [] # nothing dispatched + + +def test_guardrail_allows_a_well_separated_target(): + c = _coordinator(min_separation_m=5.0) + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 20, 0, 5) + out = c.goto_drone("drone1", north=5.0, east=0.0, altitude=5.0) + assert "REJECTED" not in out + assert c.swarm_cmd.published[-1]["action"] == "goto_ned" + assert c.swarm_cmd.published[-1]["target"] == "drone1" + + +def test_guardrail_fails_closed_without_telemetry(): + """"Cannot verify" is not "safe". No position -> refuse, do not hope.""" + c = _coordinator(min_separation_m=5.0) + _report(c, "drone1", 0, 0, 5) + reason = c._violates_separation("unknown-drone", 0.0, 0.0, -5.0) + assert reason is not None and "cannot verify" in reason + + +def test_guardrail_refuses_to_use_stale_positions(): + """At 5 m/s a stale fix is metres of travel; deciding separation on it is + arithmetic theatre, so the guardrail must refuse rather than answer.""" + from dimos.robot.drone.px4_swarm_coordinator import GUARDRAIL_MAX_AGE_SEC + + c = _coordinator(min_separation_m=5.0) + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 20, 0, 5) + c._states["drone1"]["ts"] = time.time() - (GUARDRAIL_MAX_AGE_SEC + 0.5) + reason = c._violates_separation("drone1", 19.0, 0.0, -5.0) + assert reason is not None and "cannot verify" in reason + + +def test_guardrail_still_answers_on_fresh_positions(): + c = _coordinator(min_separation_m=5.0) + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 20, 0, 5) + assert c._violates_separation("drone1", 5.0, 0.0, -5.0) is None # clear + assert c._violates_separation("drone1", 19.0, 0.0, -5.0) is not None # conflict + + +def test_fleet_altitude_cap_rejects_before_dispatch(): + c = _coordinator(max_altitude_m=30.0) + _report(c, "drone1", 0, 0, 5) + assert "REJECTED" in c.takeoff_all(100.0) + assert c.swarm_cmd.published == [] + + +# --- aggregate commands ---------------------------------------------------- + + +def test_takeoff_all_broadcasts_once(): + """takeoff is the one aggregate that stays single-shot: it is not a safety + verb, and a drone that misses it just stays on the ground -- visible and + harmless, unlike a drone that misses an rtl.""" + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + c.takeoff_all(5.0) + assert [m["action"] for m in c.swarm_cmd.published] == ["takeoff"] + + +@pytest.mark.parametrize( + ("call", "action"), + [ + ("land_all", "land"), + ("rtl_all", "rtl"), + ("hold_all", "hold"), + ("emergency_land_all", "emergency_land"), + ("kill_all", "kill"), + ], +) +def test_safety_commands_broadcast_three_times_to_all(call: str, action: str): + """Safety verbs repeat 3x on purpose: the bus is best-effort, and one lost + rtl_all broadcast left an armed drone hovering while its twin landed. All + of these are idempotent, so repetition is free insurance.""" + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 10, 0, 5) + # Make both fakes look ALREADY compliant so the (synchronous, separately + # tested) delivery verification stays quiet and this counts only the + # broadcast itself. + compliant_mode = {"land": "AUTO.LAND", "rtl": "AUTO.RTL", "hold": "AUTO.LOITER", + "emergency_land": "AUTO.LAND", "kill": None}[action] + with c._lock: + for k in ("drone1", "drone2"): + c._states[k]["mode"] = compliant_mode + c._states[k]["armed"] = False + getattr(c, call)() + assert len(c.swarm_cmd.published) == 3 + seqs = [m["seq"] for m in c.swarm_cmd.published] + assert seqs == sorted(seqs) and len(set(seqs)) == 3 + for m in c.swarm_cmd.published: + assert m["target"] == "all" and m["action"] == action + + +def test_command_sequence_numbers_increase(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + c.rtl_all() + c.hold_all() + seqs = [m["seq"] for m in c.swarm_cmd.published] + assert seqs == sorted(seqs) and len(set(seqs)) == len(seqs) + + +# --- formations and sweeps ------------------------------------------------- + + +def test_line_formation_spaces_drones_at_least_min_separation_apart(): + c = _coordinator(min_separation_m=4.0) + for i in range(3): + _report(c, f"drone{i + 1}", i * 10, 0, 5) + c.line_formation(center_north=0, center_east=0, altitude=5, spacing_m=1.0, heading_deg=90.0) + targets = [ + (m["args"]["north"], m["args"]["east"]) + for m in c.swarm_cmd.published + if m["action"] == "goto_ned" + ] + assert len(targets) == 3 + easts = sorted(e for _n, e in targets) + # spacing_m=1.0 was below the guardrail, so it must have been raised. + assert all(b - a >= 4.0 for a, b in pairwise(easts)) + + +def test_grid_sweep_gives_each_drone_its_own_strip(): + c = _coordinator(min_separation_m=2.0) + for i in range(3): + _report(c, f"drone{i + 1}", 0, i * 5, 5) + c.grid_sweep(0, 0, 60, 30, altitude=8, sub_lane_spacing_m=5) + paths = {m["target"]: m["args"]["path"] for m in c.swarm_cmd.published if m["action"] == "path"} + assert set(paths) == {"drone1", "drone2", "drone3"} + # Strips must not overlap: every drone's east range is disjoint from the others. + ranges = [] + for path in paths.values(): + easts = [wp[1] for wp in path] + ranges.append((min(easts), max(easts))) + ranges.sort() + for (_lo_a, hi_a), (lo_b, _hi_b) in pairwise(ranges): + assert lo_b > hi_a, f"strips overlap: {ranges}" + + +def test_grid_sweep_strips_respect_the_separation_buffer(): + c = _coordinator(min_separation_m=6.0) + for i in range(3): + _report(c, f"drone{i + 1}", 0, i * 5, 5) + c.grid_sweep(0, 0, 60, 60, altitude=8, sub_lane_spacing_m=5) + ranges = sorted( + (min(wp[1] for wp in m["args"]["path"]), max(wp[1] for wp in m["args"]["path"])) + for m in c.swarm_cmd.published + if m["action"] == "path" + ) + for (_lo_a, hi_a), (lo_b, _hi_b) in pairwise(ranges): + assert lo_b - hi_a >= 6.0 - 1e-6 + + +def test_grid_sweep_is_a_boustrophedon_not_a_teleport(): + """Consecutive passes must reverse direction, so the drone never jumps.""" + c = _coordinator(min_separation_m=2.0) + _report(c, "drone1", 0, 0, 5) + c.grid_sweep(0, 0, 50, 20, altitude=8, sub_lane_spacing_m=5) + path = next(m["args"]["path"] for m in c.swarm_cmd.published if m["action"] == "path") + # Waypoints come in (start, end) pairs per lane; each pair flips direction. + for i in range(0, len(path) - 2, 2): + this_dir = path[i + 1][0] - path[i][0] + next_dir = path[i + 3][0] - path[i + 2][0] + assert this_dir * next_dir < 0, "consecutive lanes must run opposite ways" + + +def test_grid_sweep_rejects_an_area_too_narrow_for_the_buffer(): + c = _coordinator(min_separation_m=10.0) + for i in range(3): + _report(c, f"drone{i + 1}", 0, i * 5, 5) + out = c.grid_sweep(0, 0, 30, 6, altitude=8) # 2 m strips, 10 m buffer needed + assert "REJECTED" in out + assert c.swarm_cmd.published == [] + + +def test_grid_sweep_altitude_is_applied_to_every_waypoint(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + c.grid_sweep(0, 0, 40, 20, altitude=12) + path = next(m["args"]["path"] for m in c.swarm_cmd.published if m["action"] == "path") + assert all(wp[2] == pytest.approx(-12.0) for wp in path) # NED down is negative altitude + + +# --- investigate ----------------------------------------------------------- + + +def test_investigate_sends_exactly_two_units_by_default(): + """The PDF's "Task two units to investigate this coordinate".""" + c = _coordinator() + for i in range(3): + _report(c, f"drone{i + 1}", i * 30, 0, 5) + c.investigate(north=100, east=50, altitude=6, num_drones=2) + sent = [m for m in c.swarm_cmd.published if m["action"] == "goto_ned"] + assert len(sent) == 2 + + +def test_investigate_staggers_targets_beyond_min_separation(): + c = _coordinator(min_separation_m=4.0) + for i in range(2): + _report(c, f"drone{i + 1}", i * 40, 0, 5) + c.investigate(north=200, east=0, altitude=6, num_drones=2) + norths = sorted(m["args"]["north"] for m in c.swarm_cmd.published if m["action"] == "goto_ned") + assert norths[1] - norths[0] >= 4.0 + + +def test_investigate_honours_an_explicit_drone_list(): + c = _coordinator() + for i in range(3): + _report(c, f"drone{i + 1}", i * 30, 0, 5) + c.investigate(north=100, east=0, drones="drone1,drone3") + targets = {m["target"] for m in c.swarm_cmd.published if m["action"] == "goto_ned"} + assert targets == {"drone1", "drone3"} + + +def test_investigate_rejects_unknown_drone_names(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + out = c.investigate(north=10, east=0, drones="drone1,ghost") + assert "Unknown drone" in out and "ghost" in out + assert c.swarm_cmd.published == [] + + +def test_fleet_skills_report_clearly_when_no_drone_has_checked_in(): + c = _coordinator() + assert "No aircraft" in c.grid_sweep(0, 0, 30, 30) + assert "No aircraft" in c.line_formation() + assert "No aircraft" in c.investigate(north=1, east=1) + assert "No robots have reported yet" in c.list_drones() + + +# --------------------------------------------------------------------------- +# Mixed fleet: air maneuvers must not be handed to ground robots +# --------------------------------------------------------------------------- + + +def _report_ground(c: SwarmCoordinator, key: str, n: float, e: float) -> None: + """Feed one legged-robot state message in.""" + lat, lon = offset_latlon(HOME_LAT, HOME_LON, n, e) + c._on_drone_state( + String( + json.dumps( + { + "key": key, + "robot_class": "legged", + "connected": True, + "ned": [n, e, -0.4], + "altitude_m": 0.4, + "battery_pct": 100.0, + "global": {"lat": lat, "lon": lon, "rel_alt_m": 0.4}, + "ts": time.time(), + } + ) + ) + ) + + +def test_grid_sweep_does_not_hand_an_air_lane_to_a_ground_robot(): + """A quadruped assigned a sweep strip silently loses that share of coverage.""" + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 10, 0, 5) + _report_ground(c, "dog1", 0, 20) + c.grid_sweep(0, 0, 60, 40, altitude=8) + targets = {m["target"] for m in c.swarm_cmd.published if m["action"] == "path"} + assert targets == {"drone1", "drone2"}, f"ground robot got an air lane: {targets}" + + +def test_investigate_counts_only_aircraft_towards_num_drones(): + """num_drones=2 must mean two *aircraft*, not one aircraft and a dog.""" + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 40, 0, 5) + _report_ground(c, "dog1", 0, 20) + c.investigate(north=100, east=0, altitude=6, num_drones=2) + sent = {m["target"] for m in c.swarm_cmd.published if m["action"] == "goto_ned"} + assert sent == {"drone1", "drone2"} + + +def test_line_formation_excludes_ground_robots(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report_ground(c, "dog1", 0, 20) + out = c.line_formation(altitude=6) + targets = {m["target"] for m in c.swarm_cmd.published if m["action"] == "goto_ned"} + assert targets == {"drone1"} + assert "dog1" in out # still reported as skipped, not silently dropped + + +def test_goto_drone_rejects_a_ground_robot_with_a_clear_reason(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report_ground(c, "dog1", 0, 20) + out = c.goto_drone("dog1", north=10, east=0, altitude=6) + assert "ground robot" in out + assert c.swarm_cmd.published == [] + + +def test_aggregate_commands_still_reach_ground_robots(): + """rtl/land are broadcast verbs every robot class translates for itself.""" + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report_ground(c, "dog1", 0, 20) + out = c.rtl_all() + assert "dog1" in out and "drone1" in out + assert c.swarm_cmd.published[-1]["target"] == "all" + + +def test_fleet_view_reports_both_classes(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report_ground(c, "dog1", 0, 20) + listing = c.list_drones() + assert "1 aircraft" in listing and "1 ground" in listing + # count_within is world-frame geometry and must span classes. + assert "dog1" in c.count_within("drone1", 100.0) + + +# --------------------------------------------------------------------------- +# Geometry helpers +# --------------------------------------------------------------------------- + + +def test_offset_latlon_round_trips_through_distance(): + lat, lon = offset_latlon(HOME_LAT, HOME_LON, 100.0, 0.0) + assert distance_3d_m((HOME_LAT, HOME_LON, 0), (lat, lon, 0)) == pytest.approx(100.0, abs=0.5) + + +def test_distance_3d_includes_altitude(): + a = (HOME_LAT, HOME_LON, 0.0) + b = (HOME_LAT, HOME_LON, 30.0) + assert distance_3d_m(a, b) == pytest.approx(30.0, abs=0.01) + + +def test_pairwise_distances_are_symmetric_and_complete(): + lat2, lon2 = offset_latlon(HOME_LAT, HOME_LON, 10, 0) + lat3, lon3 = offset_latlon(HOME_LAT, HOME_LON, 0, 10) + pairs = pairwise_distances( + {"a": (HOME_LAT, HOME_LON, 0), "b": (lat2, lon2, 0), "c": (lat3, lon3, 0)} + ) + assert len(pairs) == 3 # 3 choose 2 + assert {(a, b) for a, b, _ in pairs} == {("a", "b"), ("a", "c"), ("b", "c")} + + +# --------------------------------------------------------------------------- +# Preflight check (the PDF's hardware safety checklist, as a gate) +# --------------------------------------------------------------------------- + + +def test_preflight_passes_on_a_healthy_fleet(): + c = _coordinator(min_separation_m=2.0) + for i in range(3): + _report(c, f"drone{i + 1}", i * 10, 0, 0, sys_id=i + 1, armed=False) + assert "PASS" in c.preflight_check() + + +def test_preflight_fails_with_no_drones(): + assert "FAIL" in _coordinator().preflight_check() + + +def test_preflight_catches_duplicate_sys_ids(): + """The most common real-hardware mistake: every Pixhawk ships as MAV_SYS_ID 1.""" + c = _coordinator() + _report(c, "drone1", 0, 0, 0, sys_id=1, armed=False) + _report(c, "drone2", 20, 0, 0, sys_id=1, armed=False) + out = c.preflight_check() + assert "FAIL" in out and "MAV_SYS_ID 1 shared by" in out + + +def test_preflight_catches_low_battery(): + c = _coordinator(min_battery_pct=40.0) + _report(c, "drone1", 0, 0, 0, sys_id=1, armed=False, battery_pct=12.0) + out = c.preflight_check() + assert "FAIL" in out and "battery 12% below floor" in out + + +def test_preflight_catches_a_drone_already_armed(): + c = _coordinator() + _report(c, "drone1", 0, 0, 0, sys_id=1, armed=True) + out = c.preflight_check() + assert "FAIL" in out and "already ARMED" in out + + +def test_preflight_catches_drones_parked_too_close(): + c = _coordinator(min_separation_m=10.0) + _report(c, "drone1", 0, 0, 0, sys_id=1, armed=False) + _report(c, "drone2", 3, 0, 0, sys_id=2, armed=False) + out = c.preflight_check() + assert "FAIL" in out and "below the 10.0m minimum" in out + + +def test_preflight_catches_missing_gps_fix(): + c = _coordinator() + _report(c, "drone1", 0, 0, 0, sys_id=1, armed=False) + c._states["drone1"]["global"] = None + out = c.preflight_check() + assert "FAIL" in out and "no global position fix" in out + + +def test_preflight_catches_stale_telemetry(): + c = _coordinator() + _report(c, "drone1", 0, 0, 0, sys_id=1, armed=False) + c._states["drone1"]["ts"] = time.time() - (STATE_STALE_AFTER_SEC + 30) + out = c.preflight_check() + assert "FAIL" in out and "telemetry stale" in out + + +def test_preflight_reports_expected_but_absent_drones(): + c = _coordinator(expected="drone1,drone2,drone3") + _report(c, "drone1", 0, 0, 0, sys_id=1, armed=False) + out = c.preflight_check() + assert "drone2: expected but never reported" in out + assert "drone3: expected but never reported" in out + + +def test_preflight_always_lists_the_manual_checks(): + """Things software cannot verify must still be put in front of the operator.""" + c = _coordinator() + _report(c, "drone1", 0, 0, 0, sys_id=1, armed=False) + out = c.preflight_check() + assert "RC transmitter bound" in out + assert "Geofence configured" in out + + +def test_preflight_never_commands_anything(): + c = _coordinator() + _report(c, "drone1", 0, 0, 0, sys_id=1, armed=False) + c.preflight_check() + assert c.swarm_cmd.published == [] + + +def test_position_maneuvers_refuse_when_nothing_is_airborne(): + """A grid_sweep on grounded drones engages OFFBOARD and trips a failsafe that + then blocks the next arm, so it must be refused rather than dispatched.""" + c = _coordinator() + _report(c, "drone1", 0, 0, 0, armed=False) + _report(c, "drone2", 10, 0, 0, armed=False) + for out in ( + c.grid_sweep(0, 0, 40, 30), + c.line_formation(), + c.investigate(north=10, east=0), + c.goto_drone("drone1", north=5, east=0, altitude=5), + ): + assert "REJECTED" in out and "takeoff_all first" in out + assert c.swarm_cmd.published == [] + + +def test_position_maneuvers_use_only_airborne_aircraft(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 20, 0, 5) + _report(c, "drone3", 40, 0, 0, armed=False) # still on the ground + c.grid_sweep(0, 0, 40, 30, altitude=8) + targets = {m["target"] for m in c.swarm_cmd.published if m["action"] == "path"} + assert targets == {"drone1", "drone2"} + + +def test_preflight_does_not_fail_on_closely_parked_ground_robots(): + """The separation floor is a flight limit; two dogs standing a metre apart + is normal and must not FAIL preflight for the whole fleet.""" + c = _coordinator(min_separation_m=2.0) + _report(c, "drone1", 0, 0, 0, armed=False) + _report_ground(c, "dog1", -4.0, 0.0) + _report_ground(c, "dog2", -5.5, 0.0) # 1.5 m from dog1 + out = c.preflight_check() + assert "PASS" in out, out + assert "not a flight limit" in out + + +def test_preflight_still_fails_on_close_aircraft(): + c = _coordinator(min_separation_m=5.0) + _report(c, "drone1", 0, 0, 0, armed=False) + _report(c, "drone2", 1.0, 0, 0, armed=False) + out = c.preflight_check() + assert "FAIL" in out and "below the" in out + + +# --------------------------------------------------------------------------- +# Operating-radius boundary +# --------------------------------------------------------------------------- + + +def test_goto_drone_rejects_a_target_beyond_the_operating_radius(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + out = c.goto_drone("drone1", north=300.0, east=0.0, altitude=6.0) + assert "REJECTED" in out and "250" in out + assert c.swarm_cmd.published == [] + + +def test_goto_drone_accepts_a_target_inside_the_radius(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + out = c.goto_drone("drone1", north=249.0, east=0.0, altitude=6.0) + assert "REJECTED" not in out + assert len(c.swarm_cmd.published) == 1 + + +def test_grid_sweep_rejects_when_the_far_corner_is_outside_the_fence(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + out = c.grid_sweep(corner_b_north=400.0, corner_b_east=100.0) + assert "REJECTED" in out and "far corner" in out + assert c.swarm_cmd.published == [] + + +def test_investigate_checks_the_staggered_extreme_not_the_raw_target(): + """Two drones staggered around a target near the fence: the outermost + assignment crosses it even though the target itself does not.""" + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 0, 10, 5) + out = c.investigate(north=249.0, east=0.0, num_drones=2) + assert "REJECTED" in out + assert c.swarm_cmd.published == [] + + +def test_boundaries_reports_every_active_limit(): + c = _coordinator(max_altitude_m=40.0) + out = c.boundaries() + assert "250 m" in out + assert "40 m" in out + assert "2.0 m" in out + assert "geofence" in out + + +# --------------------------------------------------------------------------- +# Staged sweep-then-ground mission +# --------------------------------------------------------------------------- + + +def test_sweep_then_ground_needs_a_ground_robot(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + out = c.sweep_then_ground() + assert "no ground robot" in out + + +def test_sweep_then_ground_inherits_grid_sweep_gates(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report_ground(c, "dog1", -4, 0) + out = c.sweep_then_ground(corner_b_north=400.0, corner_b_east=100.0) + assert "REJECTED" in out and "far corner" in out + assert c.swarm_cmd.published == [] + + +def test_sweep_then_ground_dispatches_air_now_and_ground_after_delay(): + c = _coordinator() + c._ground_wait_timeout_s = 0.05 # dog never "arrives"; move on fast + _report(c, "drone1", 0, 0, 5) + _report_ground(c, "dog1", -4, 0) + out = c.sweep_then_ground(ground_delay_s=0.0) + assert "Staged sweep started" in out + # The air sweep is on the bus immediately (published messages arrive parsed). + assert "path" in [m["action"] for m in c.swarm_cmd.published] + # The ground legs appear as the thread times through its waypoints. + deadline = time.time() + 5.0 + while time.time() < deadline: + if [m["action"] for m in c.swarm_cmd.published].count("ground_goto") >= 4: + break + time.sleep(0.05) + gg = [m for m in c.swarm_cmd.published if m["action"] == "ground_goto"] + assert len(gg) == 4 + assert gg[0]["target"] == "dog1" + assert [g["args"]["north"] for g in gg][-1] == -4.0 # last leg is home + assert "complete" in c.mission_status() or "transect" in c.mission_status() + + +def test_emergency_aborts_the_ground_stage(): + c = _coordinator() + c._ground_wait_timeout_s = 30.0 + _report(c, "drone1", 0, 0, 5) + _report_ground(c, "dog1", -4, 0) + c.sweep_then_ground(ground_delay_s=30.0) # long delay: still in air stage + c.emergency_land_all() + deadline = time.time() + 3.0 + while time.time() < deadline and "aborted" not in c.mission_status(): + time.sleep(0.05) + assert "aborted" in c.mission_status() + + +def test_legged_module_parses_ground_goto_from_the_coordinator_wire_format(): + """The bug this pins: the handler read north/east from the TOP level of the + message, but _send() nests kwargs under "args" -- so the dog walked to the + (0, 0) defaults instead of its waypoint and the staged mission hung. The + coordinator-side tests asserted the message shape; nothing checked the + consumer parsed it. This does, end to end across the wire format.""" + from unittest.mock import MagicMock + + from dimos.robot.legged.legged_sim_module import LeggedSimModule + + # Build the exact payload _send() produces. + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report_ground(c, "dog1", -4, 0) + c._send("dog1", "ground_goto", north=12.5, east=-7.0) + wire = json.dumps(c.swarm_cmd.published[-1]) + + dog = object.__new__(LeggedSimModule) + # robot_key is a property over config.instance_name + dog.config = SimpleNamespace(instance_name="dog1") + dog.goto = MagicMock(return_value="ok") + dog.stand = MagicMock() + dog.halt = MagicMock() + dog.crouch = MagicMock() + + class _Msg: + data = wire + + dog._on_swarm_cmd(_Msg()) + dog.goto.assert_called_once_with(north=12.5, east=-7.0) + + +def test_legged_move_is_relative_to_the_robots_live_pose_and_heading(): + """Pins the frame math AND the reason `move` exists: an external caller + computing 'forward 2 m' from a stale position sends the robot back to + where it used to be. `_move_target` reads live state at call time; `move` + executes it with the heading HELD (goto would turn the robot, silently + redefining "forward" for the next command).""" + from unittest.mock import MagicMock + + from dimos.robot.legged.legged_sim_module import LeggedSimModule + + dog = object.__new__(LeggedSimModule) + dog.config = SimpleNamespace(instance_name="dog1", max_range_m=250.0) + + def set_pose(n, e, yaw): + dog._state_dict = MagicMock( + return_value={"connected": True, "ned": [n, e, 0.0], "yaw_ned_rad": yaw} + ) + + # Facing NORTH at (10, 5): forward -> +north, left 5 (right_m=-5) -> WEST. + set_pose(10.0, 5.0, 0.0) + tn, te, yaw0 = dog._move_target(2.0, -5.0) + assert abs(tn - 12.0) < 1e-6 and abs(te - 0.0) < 1e-6 and yaw0 == 0.0 + + # Facing EAST (yaw pi/2): forward -> +east, right -> SOUTH. + set_pose(0.0, 0.0, math.pi / 2) + tn, te, _ = dog._move_target(3.0, 1.0) + assert abs(tn - (-1.0)) < 1e-6 and abs(te - 3.0) < 1e-6 + + # The chaining property the user's bug was about: second move computes + # from the NEW pose, not the original one. + set_pose(0.0, -5.0, 0.0) # after "left 5" + tn, te, _ = dog._move_target(2.0, 0.0) + assert abs(tn - 2.0) < 1e-6 and abs(te - (-5.0)) < 1e-6 + + +def test_legged_move_rejects_targets_beyond_the_operating_radius(): + from unittest.mock import MagicMock + + from dimos.robot.legged.legged_sim_module import LeggedSimModule + + dog = object.__new__(LeggedSimModule) + dog.config = SimpleNamespace(instance_name="dog1", max_range_m=250.0) + dog._state_dict = MagicMock( + return_value={"connected": True, "ned": [249.0, 0.0, 0.0], "yaw_ned_rad": 0.0} + ) + out = dog.move(forward_m=10.0) + assert "REJECTED" in out and "250" in out + + +def test_bridge_move_controller_math(): + """Pins _move_cmds: frame conversions and the deadband bumps, in the NWU + bridge frame with NED-signed outputs.""" + from dimos.simulation.px4_hil.fleet_bridge import _move_cmds + + # Target dead ahead (north), facing north: pure forward, no strafe/turn. + vx, vy, wz = _move_cmds(5.0, 0.0, 0.0, 0.0) + assert vx == 0.7 and vy == 0.0 and wz == 0.0 + + # Target to the WEST (NWU +y), facing north: west is LEFT -> vy negative. + vx, vy, wz = _move_cmds(0.0, 5.0, 0.0, 0.0) + assert vx == 0.0 and vy == -0.3 and wz == 0.0 + + # Heading disturbed LEFT of hold (NWU yaw +0.2): correction turns RIGHT + # (NED wz positive). + _, _, wz = _move_cmds(0.0, 0.0, 0.2, 0.0) + assert wz > 0.0 + + # Deadband bump: small forward error still above half the arrive radius + # commands at least 0.3, never a sub-deadband creep. + vx, _, _ = _move_cmds(0.30, 0.0, 0.0, 0.0) + assert abs(vx) >= 0.3 + + +def test_drone_aero_helpers(): + """Pins the aero layer's pure math: spool lag, ground effect, H-drag.""" + import numpy as np + + from dimos.simulation.px4_hil.hil_bridge import ( + GROUND_EFFECT_MAX, + _ground_effect_boost, + _motor_lag_step, + _rotor_hdrag_force, + ) + + # Spool: converges toward the command, never overshoots, monotone. + st = np.zeros(4) + cmd = np.full(4, 0.8) + prev = st.copy() + for _ in range(200): + st = _motor_lag_step(st, cmd, 0.004) + assert np.all(st >= prev - 1e-12) and np.all(st <= 0.8 + 1e-12) + prev = st.copy() + assert np.all(st > 0.79) # settled after ~0.8 s at tau=0.06 + + # Ground effect: capped in the clamped near-ground region, present but + # modest at half a metre, gone at altitude, monotonically decreasing. + assert _ground_effect_boost(0.05) == GROUND_EFFECT_MAX + assert 0.0 < _ground_effect_boost(0.5) < _ground_effect_boost(0.2) < GROUND_EFFECT_MAX + assert _ground_effect_boost(10.0) < 0.001 + + # H-drag opposes airspeed, scales with thrust, no vertical component. + f = _rotor_hdrag_force(np.array([2.0, -1.0, 3.0]), 20.0) + assert f[0] < 0 and f[1] > 0 and f[2] == 0.0 + f2 = _rotor_hdrag_force(np.array([2.0, -1.0, 0.0]), 40.0) + assert abs(f2[0]) > abs(f[0]) + # Motors off, falling: no phantom damping. + assert np.allclose(_rotor_hdrag_force(np.array([5.0, 5.0, -8.0]), 0.0), 0.0) + + +def test_fleet_bridge_constructs_without_vehicles(): + """Pins the whole Px4HilFleet.__init__ path (wind setup included) without + needing PX4 or sockets. A NameError here shipped once: the wind code used + `np` that was only imported inside Go1Link, so EVERY sim start died at + construction while all unit tests stayed green -- none of them built the + fleet.""" + from dimos.simulation.px4_hil.fleet_bridge import Px4HilFleet + + fleet = Px4HilFleet(n_drones=0, n_dogs=0) + assert fleet.links == [] and fleet.dogs == [] + assert fleet._wind_mean.shape == (3,) + + +def test_drone_link_apply_runs_the_full_aero_path(): + """Constructs a real DroneLink on a real model (no sockets, no PX4) and + runs apply() with wind -- the path where a missing import shipped TWICE + (np in the fleet ctor, MAX_THRUST_PER_ROTOR_N in apply) while the pure- + math helper tests stayed green. Steps the world to prove forces are + finite and the drone does not explode at rest.""" + import mujoco + import numpy as np + + from dimos.simulation.px4_hil.hil_bridge import DroneLink + from dimos.simulation.px4_hil.scene import build_model + + model, _ = build_model(1, 0) + data = mujoco.MjData(model) + mujoco.mj_forward(model, data) + link = DroneLink(mujoco, model, data, 0, "127.0.0.1", 4560) + link.controls[:] = 0.5 + wind = np.array([3.0, 0.0, 0.0]) + for _ in range(50): + link.apply(wind) + mujoco.mj_step(model, data) + assert np.all(np.isfinite(data.qpos)) and np.all(np.isfinite(data.qvel)) + assert np.all(link.motor_state > 0.1) # spool actually progressed + + +def test_fleet_verify_resends_to_the_straggler_only(): + """The stranded-drone incident, mechanized: after rtl_all, a drone still + reporting AUTO.LOITER gets an individually addressed re-send; the one that + complied gets nothing.""" + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + _report(c, "drone2", 10, 0, 5) + with c._lock: + c._states["drone1"]["mode"] = "AUTO.LOITER" # never heard the rtl + c._states["drone2"]["mode"] = "AUTO.RTL" # complied + c._verify_rounds = 1 + c._verify_fleet_command("rtl", ["drone1", "drone2"]) + targeted = [m for m in c.swarm_cmd.published if m["target"] == "drone1"] + assert targeted and all(m["action"] == "rtl" for m in targeted) + assert not any(m["target"] == "drone2" for m in c.swarm_cmd.published) + + +def test_fleet_verify_is_quiet_when_everyone_complied(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + with c._lock: + c._states["drone1"]["mode"] = "AUTO.RTL" + c._verify_fleet_command("rtl", ["drone1"]) + assert c.swarm_cmd.published == [] + + +def test_fleet_verify_kill_checks_armed_not_mode(): + c = _coordinator() + _report(c, "drone1", 0, 0, 5) + with c._lock: + c._states["drone1"]["armed"] = True # kill did not land + c._verify_rounds = 1 + c._verify_fleet_command("kill", ["drone1"]) + assert any( + m["target"] == "drone1" and m["action"] == "kill" for m in c.swarm_cmd.published + ) diff --git a/dimos/robot/drone/test_px4_swarm_bus.py b/dimos/robot/drone/test_px4_swarm_bus.py new file mode 100644 index 0000000000..e54918c6f7 --- /dev/null +++ b/dimos/robot/drone/test_px4_swarm_bus.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration test for the swarm fleet bus across a namespace boundary. + +The PX4 swarm depends on two *exposed* streams doing real work at runtime: + +* ``drone_state`` — N namespaced drones publish, one shared coordinator listens + (fan-in, N->1). +* ``swarm_cmd`` — one shared coordinator publishes, N namespaced drones listen + (fan-out, 1->N). + +Unit tests can only prove the blueprint *wiring* is right. This one actually +stands a coordinator up and checks messages traverse the namespace boundary in +both directions, because a silent failure here looks exactly like "the drones +ignored the command". + +Stand-in modules are used rather than the real Px4DroneModule/SwarmCoordinator +so the test needs no MAVLink link and no running sim. +""" + +from __future__ import annotations + +import time +from typing import Any + +from dimos_lcm.std_msgs import String +import pytest + +from dimos.core.coordination.blueprint_config.parser import BlueprintConfigParser +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.coordination.module_coordinator import ModuleCoordinator +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.stream import In, Out + +# The two streams the swarm leaves unprefixed so they cross the boundary. +BUS = {"swarm_cmd", "drone_state"} + +FLEET_SIZE = 3 +SETTLE_SEC = 2.0 +DELIVERY_TIMEOUT_SEC = 15.0 + + +class BusLeaf(Module): + """Stand-in for Px4DroneModule: listens on swarm_cmd, publishes drone_state.""" + + swarm_cmd: In[String] + drone_state: Out[String] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._got: list[str] = [] + + @rpc + def start(self) -> None: + super().start() + if getattr(self.swarm_cmd, "transport", None): + self.swarm_cmd.subscribe(self._on_cmd) + + def _on_cmd(self, msg: String) -> None: + self._got.append(msg.data) + + @rpc + def received(self) -> list[str]: + """Read back over RPC — the callback runs in this module's worker process.""" + return list(self._got) + + @rpc + def swarm_cmd_has_transport(self) -> bool: + return getattr(self.swarm_cmd, "_transport", None) is not None + + @rpc + def report_state(self, text: str) -> str: + self.drone_state.publish(String(text)) + return "published" + + +class BusHub(Module): + """Stand-in for SwarmCoordinator: publishes swarm_cmd, listens on drone_state.""" + + drone_state: In[String] + swarm_cmd: Out[String] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._got: list[str] = [] + + @rpc + def start(self) -> None: + super().start() + if getattr(self.drone_state, "transport", None): + self.drone_state.subscribe(self._on_state) + + def _on_state(self, msg: String) -> None: + self._got.append(msg.data) + + @rpc + def received(self) -> list[str]: + return list(self._got) + + @rpc + def swarm_cmd_has_transport(self) -> bool: + return getattr(self.swarm_cmd, "_transport", None) is not None + + @rpc + def broadcast(self, text: str) -> str: + self.swarm_cmd.publish(String(text)) + return "published" + + +def _swarm_bus_blueprint(): + return autoconnect( + BusHub.blueprint(), + *[ + BusLeaf.blueprint().namespace(f"drone{i + 1}", expose=BUS) + for i in range(FLEET_SIZE) + ], + ) + + +def _wait_for(predicate, timeout: float = DELIVERY_TIMEOUT_SEC) -> bool: + """Poll until predicate() is truthy. Pub/sub delivery is asynchronous.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.25) + return False + + +@pytest.fixture +def bus_coordinator(): + blueprint = _swarm_bus_blueprint() + parsed = BlueprintConfigParser(blueprint).parse( + environ={}, overrides={"g": {"viewer": "none"}} + ) + coordinator = ModuleCoordinator.build(blueprint, parsed) + # Let every module's start() run and its subscriptions land before publishing. + time.sleep(SETTLE_SEC) + try: + yield coordinator + finally: + coordinator.stop() + + +def test_exposed_out_is_actually_bound_to_a_transport(bus_coordinator): + """A silently-unbound Out would make every fleet command a no-op.""" + hub = bus_coordinator.get_instance("bushub") + assert hub.swarm_cmd_has_transport() is True + for i in range(FLEET_SIZE): + leaf = bus_coordinator.get_instance(f"drone{i + 1}/busleaf") + assert leaf.swarm_cmd_has_transport() is True + + +def test_shared_publisher_fans_out_to_every_namespaced_listener(bus_coordinator): + """swarm_cmd: 1 -> N across the namespace boundary (takeoff_all, rtl_all, ...).""" + hub = bus_coordinator.get_instance("bushub") + hub.broadcast("takeoff-all") + + leaves = [bus_coordinator.get_instance(f"drone{i + 1}/busleaf") for i in range(FLEET_SIZE)] + assert _wait_for(lambda: all("takeoff-all" in leaf.received() for leaf in leaves)), ( + "swarm_cmd did not reach every drone: " + + repr([leaf.received() for leaf in leaves]) + ) + + +def test_namespaced_publishers_fan_in_to_the_shared_listener(bus_coordinator): + """drone_state: N -> 1 across the namespace boundary (fleet_state).""" + for i in range(FLEET_SIZE): + bus_coordinator.get_instance(f"drone{i + 1}/busleaf").report_state(f"state-{i + 1}") + + hub = bus_coordinator.get_instance("bushub") + expected = {f"state-{i + 1}" for i in range(FLEET_SIZE)} + assert _wait_for(lambda: expected <= set(hub.received())), ( + f"coordinator did not see every drone's state: {hub.received()!r}" + ) From a57cd0285b30644e7a60e35bef2966888ee3a3f8 Mon Sep 17 00:00:00 2001 From: Ezenbaatar Batjargal Date: Mon, 31 Aug 2026 11:38:36 -0700 Subject: [PATCH 3/4] feat(legged): simulated quadruped module and mixed-fleet blueprints LeggedSimModule mirrors the drone module on the shared fleet bus (same state shape, robot_class field) so coordinator queries span embodiments. Real Menagerie Go1 driven by the repo's trained ONNX policy when assets are present, primitive fallback otherwise; heading-held relative move runs inside the simulator at physics rate; ground fence and inter-robot proximity halt. mixed-fleet-mcp and mixed-fleet-agentic compose N drones + M dogs from SIM_DRONES/SIM_DOGS. Co-Authored-By: Claude Fable 5 --- dimos/robot/all_blueprints.py | 3 + .../blueprints/agentic/mixed_fleet_agentic.py | 131 ++++ .../blueprints/basic/mixed_fleet_mcp.py | 128 ++++ dimos/robot/legged/legged_sim_module.py | 609 ++++++++++++++++++ 4 files changed, 871 insertions(+) create mode 100644 dimos/robot/legged/blueprints/agentic/mixed_fleet_agentic.py create mode 100644 dimos/robot/legged/blueprints/basic/mixed_fleet_mcp.py create mode 100644 dimos/robot/legged/legged_sim_module.py diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 881abbba18..e59cbc0351 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -121,6 +121,9 @@ "unitree-g1-record": "dimos.robot.unitree.g1.blueprints.basic.unitree_g1_record:unitree_g1_record", "unitree-g1-shm": "dimos.robot.unitree.g1.blueprints.perceptive.unitree_g1_shm:unitree_g1_shm", "unitree-g1-sim": "dimos.robot.unitree.g1.blueprints.perceptive.unitree_g1_sim:unitree_g1_sim", + "drone-px4-sitl-fleet-mcp": "dimos.robot.drone.blueprints.basic.drone_px4_sitl_fleet_mcp:drone_px4_sitl_fleet_mcp", + "mixed-fleet-agentic": "dimos.robot.legged.blueprints.agentic.mixed_fleet_agentic:mixed_fleet_agentic", + "mixed-fleet-mcp": "dimos.robot.legged.blueprints.basic.mixed_fleet_mcp:mixed_fleet_mcp", "unitree-go2": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2", "unitree-go2-agentic": "dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_agentic:unitree_go2_agentic", "unitree-go2-agentic-huggingface": "dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_agentic_huggingface:unitree_go2_agentic_huggingface", diff --git a/dimos/robot/legged/blueprints/agentic/mixed_fleet_agentic.py b/dimos/robot/legged/blueprints/agentic/mixed_fleet_agentic.py new file mode 100644 index 0000000000..36098af6a4 --- /dev/null +++ b/dimos/robot/legged/blueprints/agentic/mixed_fleet_agentic.py @@ -0,0 +1,131 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Natural-language control of the mixed drone + legged fleet. + +Same fleet as ``mixed-fleet-mcp``, plus an agent, so the whole thing can be +driven in plain English instead of one ``dimos mcp call`` at a time: + + ./dimos/simulation/px4_hil/sim.sh start 2 1 --viewer + OPENAI_API_KEY=... CI=1 SIM_DRONES=2 SIM_DOGS=1 \\ + dimos run mixed-fleet-agentic + dimos humancli + + > take both drones up to 6 metres and sweep a 40 by 30 metre area + > walk the dog 10 metres north, then bring the drones home + +``dimos mcp call`` still works exactly as before -- the agent sits alongside the +MCP server, it does not replace it. +""" + +from dimos.agents.mcp.mcp_client import McpClient +from dimos.agents.mcp.mcp_server import McpServer +from dimos.core.coordination.blueprints import autoconnect +from dimos.robot.legged.blueprints.basic.mixed_fleet_mcp import ( + _DRONES, + _N_DOGS, + mixed_fleet, +) + +MIXED_FLEET_SYSTEM_PROMPT = """ +You command a mixed robot fleet in simulation: multirotor drones running real +PX4 firmware, and Unitree Go1 quadrupeds. Both share one physics world. + +TOOL NAMING +- Fleet-wide tools have plain names: list_drones, preflight_check, fleet_state, + takeoff_all, land_all, rtl_all, hold_all, grid_sweep, line_formation, + investigate, count_within, goto_drone, emergency_land_all, kill_all. +- A skill only one robot offers keeps its bare name (walk, stand, halt). +- A skill several robots offer is qualified with the robot's name and module, + joined by UNDERSCORES in your tool list: drone2_px4dronemodule_takeoff, + dog1_leggedsimmodule_goto. (Inside DimOS the separator is "/"; your tool + names use "_" because that is what the API allows -- each such tool's + description starts with its real [bracketed/slashed] name.) +- `state` and `goto` exist on BOTH robot classes, so they are always qualified. +- Use ONLY names exactly as they appear in your tool list. If unsure, call + list_drones first and match the keys. + +ORDERING RULE -- THIS ONE MATTERS +- Call takeoff_all BEFORE any position maneuver (grid_sweep, line_formation, + investigate, goto_drone). Those engage OFFBOARD mode, and commanding OFFBOARD + to a grounded, disarmed drone trips a PX4 failsafe that then blocks the NEXT + arm attempt. The failure shows up minutes later on a different command. +- Aircraft-only maneuvers ignore the ground robots. That is correct, not a bug; + say so rather than trying to make a dog fly. + +DRONES +- Altitudes are metres, positive up. 3 m is a safe default; use 5-10 m for + sweeps. Positions are metres in each drone's local NED frame, origin at its + own home. +- Prefer goto_drone over a drone's own goto while others are airborne: only the + coordinator sees the whole fleet and can reject a target that would break + minimum separation. If a target is rejected, choose a different one -- never + retry the same waypoint. +- "sweep an X metre radius" -> grid_sweep with corner_a=(-X,-X), corner_b=(X,X). +- "an X by Y area" -> grid_sweep with corner_a=(0,0), corner_b=(X,Y). +- A sweep does NOT auto-return; drones hover at their lane ends. If the user + meant "search then come back", call rtl_all afterwards. Ask if unsure. + +LEGGED ROBOTS (Unitree Go1, trained locomotion policy) +- walk(speed_mps, turn_rate_rads) drives it; halt stops it; stand settles it and + also clears a fall latch. The tool is `halt`, NOT `stop` -- `stop` is the + framework's module-teardown RPC and is not a robot command. +- Envelope: 0.94 m/s forward, 0.31 m/s reverse, 0.55 rad/s yaw. Positive turn + rate is nose-RIGHT. +- IMPORTANT: the policy has a low-speed deadband. A commanded speed below about + 0.25 m/s produces NO motion at all. Never command 0.1-0.2 m/s expecting a slow + walk -- use 0.3 m/s or more, or stop. +- crouch and set_height do nothing on a real Go1: the policy has no height + control. Tell the user rather than pretending it worked. +- RELATIVE motion ("move left 5 meters", "go forward 2", "back up"): + ALWAYS use dogN_leggedsimmodule_move(forward_m, right_m). It computes the + target from the robot's LIVE position and heading at call time, so + consecutive relative commands chain correctly. left = negative right_m, + back = negative forward_m. +- NEVER compute a relative target yourself from a position you observed + earlier and pass it to goto: by the time you act the robot has moved, and + it will walk BACK to where it used to be. That exact bug is why `move` + exists. +- ABSOLUTE motion ("go to point (10, 5)", "go to the far edge of the field"): + use dogN_leggedsimmodule_goto(north, east) -- world coordinates in metres. +- Sequential orders in one request ("left 5, then forward 2"): issue the first + move, wait for arrival (poll dogN_leggedsimmodule_state until walking is + false), then issue the next. Firing both at once makes the second cancel + the first. +- NEVER measure a distance with a timed walk (walk for N seconds then halt): + the simulation runs many times faster than your clock, so the distance is + random. Distances are ONLY move (relative) or goto (absolute). +- NEVER track the robot's position in your head across commands. Read + dogN_leggedsimmodule_state when you need a position; `move` already reads + the live pose for you. + +SAFETY +- emergency_land_all: every drone descends under its own control. +- kill_all: force-disarms everything and they FALL. Only when crashing beats + what the fleet is about to do -- and say plainly that you are doing it. +- preflight_check is read-only and never commands anything. Use it freely. + +STYLE +- Report what you did in one short summary. Do not narrate every call. +- "Dispatched" is not "done". If the user asks whether something worked, check + fleet_state rather than assuming the dispatch succeeded. +""" + +mixed_fleet_agentic = autoconnect( + mixed_fleet(), + McpServer.blueprint(), + McpClient.blueprint(system_prompt=MIXED_FLEET_SYSTEM_PROMPT, model="gpt-4o"), +).global_config(n_workers=max(4, 2 * (len(_DRONES) + _N_DOGS) + 1)) + +__all__ = ["MIXED_FLEET_SYSTEM_PROMPT", "mixed_fleet_agentic"] diff --git a/dimos/robot/legged/blueprints/basic/mixed_fleet_mcp.py b/dimos/robot/legged/blueprints/basic/mixed_fleet_mcp.py new file mode 100644 index 0000000000..8c282b6b3a --- /dev/null +++ b/dimos/robot/legged/blueprints/basic/mixed_fleet_mcp.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Mixed fleet: N namespaced PX4 drones and M namespaced legged robots, one world. + +Both robot classes are composed the same way -- one module instance per vehicle +under its own namespace, publishing onto one shared fleet bus. The coordinator +does not know or care which is which; it reads ``robot_class`` off the state +messages. That uniformity is the point of the namespace architecture, and it is +what lets an operator say "anything that flies" or "anything on the ground". + + drone1/px4dronemodule/takeoff aircraft + dog1/leggedsimmodule/stand ground robot + fleet_state both, one report + +Bring-up (the simulator must be listening before PX4 starts): + + python -m dimos.simulation.px4_hil.fleet_bridge --drones 3 --dogs 1 + # once per drone: + cd ~/PX4-Autopilot/build/px4_sitl_default && PX4_SIM_MODEL=none_iris ./bin/px4 -i 0 -d + python dimos/simulation/px4_hil/sim_params.py --count 3 + CI=1 dimos run mixed-fleet-mcp --daemon + +Then:: + + dimos mcp call fleet_state + dimos mcp call takeoff_all --arg altitude=6 + dimos mcp call dog1/leggedsimmodule/crouch + dimos mcp call rtl_all +""" + +import os + +from dimos.agents.mcp.mcp_server import McpServer +from dimos.core.coordination.blueprints import autoconnect +from dimos.robot.drone.blueprints.basic.drone_px4_sitl_fleet_mcp import FLEET_BUS +from dimos.robot.drone.px4_drone_module import Px4DroneModule +from dimos.robot.drone.px4_sitl_fleet_config import get_px4_sitl_fleet_configs +from dimos.robot.drone.px4_swarm_coordinator import SwarmCoordinator +from dimos.robot.legged.legged_sim_module import LeggedSimModule +from dimos.simulation.px4_hil.fleet_bridge import LEGGED_PORT_BASE + +# Fleet size comes from the environment so one blueprint covers every mix +# without editing code. The MuJoCo bridge is the authority on what actually +# exists in the world -- these must match what you passed to +# `fleet_bridge --drones N --dogs M`, or DimOS will wait for vehicles that +# are not there. +# +# SIM_DRONES=2 SIM_DOGS=0 dimos run mixed-fleet-mcp +DEFAULT_DRONE_COUNT = 3 +DEFAULT_DOG_COUNT = 1 + + +def _count(var: str, default: int) -> int: + raw = os.environ.get(var) + if raw is None or not raw.strip(): + return default + try: + value = int(raw) + except ValueError as e: + raise ValueError(f"{var} must be an integer, got {raw!r}") from e + if value < 0: + raise ValueError(f"{var} must be >= 0, got {value}") + return value + + +_N_DRONES = _count("SIM_DRONES", DEFAULT_DRONE_COUNT) +_N_DOGS = _count("SIM_DOGS", DEFAULT_DOG_COUNT) + +# fleet_size=0 is legitimate (a dogs-only world), but the config helper rejects +# it, so ask for nothing rather than asking for zero. +_DRONES = get_px4_sitl_fleet_configs(fleet_size=_N_DRONES) if _N_DRONES else [] + + +def mixed_fleet( + drone_configs=_DRONES, + n_dogs: int = _N_DOGS, + min_separation_m: float = 2.0, + # 40 m default: below the 45 m PX4 geofence ceiling, so a too-high command + # gets the coordinator's polite refusal, never the fence. Pass None to + # remove the fleet cap (each drone still enforces its own). + max_altitude_m: float | None = 40.0, +): + """Compose namespaced drones and legged robots around one shared coordinator.""" + drone_keys = [f"drone{i + 1}" for i in range(len(drone_configs))] + dog_keys = [f"dog{i + 1}" for i in range(n_dogs)] + return autoconnect( + SwarmCoordinator.blueprint( + min_separation_m=min_separation_m, + max_altitude_m=max_altitude_m, + expected_drones=",".join(drone_keys + dog_keys), + ), + *[ + Px4DroneModule.blueprint( + connection_string=c.connection_string, + instance=c.instance, + sys_id=c.sys_id, + max_altitude_m=max_altitude_m, + ).namespace(key, expose=FLEET_BUS) + for key, c in zip(drone_keys, drone_configs, strict=True) + ], + *[ + LeggedSimModule.blueprint( + endpoint=f"udp:127.0.0.1:{LEGGED_PORT_BASE + i}", + ).namespace(key, expose=FLEET_BUS) + for i, key in enumerate(dog_keys) + ], + ) + + +mixed_fleet_mcp = autoconnect( + mixed_fleet(), + McpServer.blueprint(), +).global_config(n_workers=max(3, 2 * (len(_DRONES) + _N_DOGS))) + +__all__ = ["DEFAULT_DOG_COUNT", "mixed_fleet", "mixed_fleet_mcp"] diff --git a/dimos/robot/legged/legged_sim_module.py b/dimos/robot/legged/legged_sim_module.py new file mode 100644 index 0000000000..e154132f93 --- /dev/null +++ b/dimos/robot/legged/legged_sim_module.py @@ -0,0 +1,609 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""One simulated legged robot, one module instance -- the ground-robot twin of +[Px4DroneModule][dimos.robot.drone.px4_drone_module.Px4DroneModule]. + +Deliberately the same shape as the drone module, because the whole point of the +namespace architecture is that a mixed fleet addresses uniformly: + + drone1/px4dronemodule/takeoff + dog1/leggedsimmodule/stand + +Each instance owns one UDP link to the MuJoCo fleet bridge, exactly as the drone +module owns one MAVLink link to a PX4 instance. Under +``.namespace("dog1", expose={"robot_state", "swarm_cmd"})`` it gets its own RPC +surface, its own topics and its own config keys, and it publishes onto the +*same* fleet bus the drones use. + +That shared bus is what makes ``SwarmCoordinator`` fleet-aware rather than +drone-aware: every robot publishes a state message carrying ``key``, +``robot_class`` and a position, so ``fleet_state`` and ``count_within`` cover +aircraft and ground robots without special-casing either. + +Positions are reported in the simulator's NWU world frame and converted to the +NED convention the rest of the fleet stack uses, so distances between a drone +and a dog are computed in one frame. +""" + +from __future__ import annotations + +import json +import math +import socket +import threading +import time +from typing import Any + +from dimos_lcm.std_msgs import String + +from dimos.agents.annotation import skill +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.robot.drone.px4_geo import offset_latlon +from dimos.simulation.px4_hil.hil_bridge import ( + DEFAULT_ORIGIN_ALT, + DEFAULT_ORIGIN_LAT, + DEFAULT_ORIGIN_LON, +) +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +# Republish cadence onto the shared fleet bus. Matches Px4DroneModule so the +# fleet-wide picture ages uniformly across robot classes. +STATE_PUBLISH_HZ = 4.0 +# The bridge only answers a peer it has heard from, so keep poking it until it +# does. Cheap, and it means module restarts reconnect on their own. +HELLO_INTERVAL_S = 1.0 + +# Closed-loop `goto` tuning. +# +# This gait turns roughly 5x better while walking than on the spot (~0.22 rad/s +# versus ~0.047), because yaw comes only from a differential stride. So the +# controller never slows down to correct heading -- it holds cruise speed and +# turns through a wide arc. Creeping while misaligned, which is the obvious +# design, actually walks *away* from the target faster than the weakened turn +# can correct, and the distance grows. +GOTO_ARRIVE_M = 0.8 +GOTO_YAW_GAIN = 0.9 # rad/s of turn per rad of heading error +# Speeds are sized for the trained Go1 policy, which has a hard deadband: a +# commanded 0.18 m/s produces NO motion at all (measured -- the robot stood +# still through a whole goto), 0.30 walks, 0.45 walks well, and 1.0 tracks at +# 0.94. Anything below ~0.25 m/s is therefore not a slow walk, it is a stop. +# The primitive-quadruped fallback clamps these down to its own lower limits, +# so one set of constants serves both backends. +GOTO_CRUISE_MS = 0.70 +# Only slow down on the final approach, where overshoot costs more than time. +GOTO_SLOWDOWN_M = 4.0 +GOTO_MIN_SPEED_MS = 0.35 +# Deadband on heading error, so the controller stops chattering the turn +# direction back and forth once it is roughly on course. +GOTO_YAW_DEADBAND_RAD = 0.12 +GOTO_TICK_S = 0.2 +GOTO_TIMEOUT_S = 300.0 + + +class LeggedSimConfig(ModuleConfig): + """Per-robot identity. Addressable as `-o dog2/leggedsimmodule.endpoint=...`.""" + + # UDP endpoint of this robot on the MuJoCo fleet bridge (LEGGED_PORT_BASE + i). + endpoint: str = "udp:127.0.0.1:15000" + # Nominal standing height, metres. Used to normalise `set_height`. + nominal_height_m: float = 0.40 + # World origin, shared with the drones. SwarmCoordinator does all its + # geometry in world frame off the `global` field, so a robot without one is + # invisible to fleet_state distances and count_within. Defaults match the + # PX4 SITL origin the HIL bridge uses. + origin_lat: float = DEFAULT_ORIGIN_LAT + origin_lon: float = DEFAULT_ORIGIN_LON + origin_alt: float = DEFAULT_ORIGIN_ALT + # Operating radius for `goto`, metres from the world origin. Matches the + # aircraft coordinator's max_range_m so one boundary governs the whole + # fleet; the sim ground plane ends at 300 m. + max_range_m: float = 250.0 + + +class LeggedSimModule(Module): + """A single simulated quadruped: stance control plus fleet-bus telemetry.""" + + config: LeggedSimConfig + + # Exposed (global): the shared fleet bus, same streams the drones use. + drone_state: Out[String] + swarm_cmd: In[String] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._sock: socket.socket | None = None + self._addr: tuple[str, int] | None = None + self._running = False + self._rx_thread: threading.Thread | None = None + self._tx_thread: threading.Thread | None = None + self._latest: dict[str, Any] = {} + self._lock = threading.RLock() + self._goto_thread: threading.Thread | None = None + self._goto_stop = threading.Event() + + # -- identity ---------------------------------------------------------- + + @property + def robot_key(self) -> str: + """Fleet-facing name, taken from the namespace prefix (`dog1/...` -> `dog1`).""" + name = self.config.instance_name + if name and "/" in name: + return name.rsplit("/", 1)[0] + return name or "dog1" + + # -- lifecycle --------------------------------------------------------- + + @rpc + def start(self) -> None: + super().start() + host, port = self._parse_endpoint(self.config.endpoint) + self._addr = (host, port) + self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + self._sock.settimeout(0.5) + self._running = True + + self._rx_thread = threading.Thread( + target=self._rx_loop, daemon=True, name=f"legged-rx-{self.robot_key}" + ) + self._rx_thread.start() + self._tx_thread = threading.Thread( + target=self._state_loop, daemon=True, name=f"legged-tx-{self.robot_key}" + ) + self._tx_thread.start() + + if getattr(self.swarm_cmd, "transport", None): + self.swarm_cmd.subscribe(self._on_swarm_cmd) + logger.info(f"[{self.robot_key}] subscribed to swarm_cmd") + logger.info(f"[{self.robot_key}] legged sim link -> {host}:{port}") + + def stop(self) -> None: + """Framework teardown. NOT the agent-facing stop -- that is ``halt``. + + Overrides ``Module.stop()``. Do not add a ``@skill`` named ``stop`` to + this class: it would shadow this method and the module would never close + its RPC, tools or event loop. + """ + self._running = False + for t in (self._rx_thread, self._tx_thread): + if t is not None and t.is_alive(): + t.join(timeout=1.0) + if self._sock is not None: + self._sock.close() + super().stop() + + @staticmethod + def _parse_endpoint(endpoint: str) -> tuple[str, int]: + raw = endpoint.split("://", 1)[-1] if "://" in endpoint else endpoint + raw = raw[4:] if raw.startswith("udp:") else raw + host, _, port = raw.rpartition(":") + return (host or "127.0.0.1", int(port)) + + # -- transport --------------------------------------------------------- + + def _send(self, payload: dict[str, Any]) -> None: + if self._sock is None or self._addr is None: + return + try: + self._sock.sendto(json.dumps(payload).encode(), self._addr) + except OSError as e: + logger.debug(f"[{self.robot_key}] send failed: {e}") + + def _rx_loop(self) -> None: + """Receive ground truth from the bridge; re-announce until it replies.""" + last_hello = 0.0 + while self._running: + now = time.monotonic() + if not self._latest and now - last_hello > HELLO_INTERVAL_S: + # The bridge learns our address from any datagram; a no-op + # stance command is the cheapest way to introduce ourselves. + self._send({"action": "noop"}) + last_hello = now + if self._sock is None: + return + try: + payload, _ = self._sock.recvfrom(4096) + except TimeoutError: + continue + except OSError: + return + try: + state = json.loads(payload) + except ValueError: + continue + with self._lock: + self._latest = state + self._enforce_ground_fence(state) + + def _enforce_ground_fence(self, state: dict[str, Any]) -> None: + """Halt a robot that walks out of the operating radius. + + The aircraft get this backstop from PX4's own geofence; a ground robot + on an open-loop ``walk`` had NOTHING -- ``goto`` refuses far targets at + dispatch, but a plain "walk forward" continues until the 300 m edge of + the world. This is the layer that catches what dispatch checks cannot, + applied at the same radius the whole fleet uses. + + Hysteresis: halts once at the fence, re-arms only after the robot is + back inside 95% of the radius, so a robot sitting on the line does not + get halt-spammed. + """ + nwu = state.get("nwu") + if not nwu: + return + # Bridge frame is NWU (y = West); NED east = -y. + r = math.hypot(float(nwu[0]), float(nwu[1])) + limit = self.config.max_range_m + breached = getattr(self, "_fence_breached", False) + if not breached: + if r > limit and state.get("walking"): + self._fence_breached = True + self._fence_halt_r = r + self._cancel_goto() + self._send({"action": "stop"}) + logger.warning( + f"[{self.robot_key}] GROUND FENCE: {r:.0f} m from origin exceeds " + f"{limit:.0f} m -- halted. Walk it back with a goto or walk command." + ) + elif r < limit * 0.95: + self._fence_breached = False + logger.info(f"[{self.robot_key}] back inside the operating radius ({r:.0f} m)") + elif state.get("walking") and r > getattr(self, "_fence_halt_r", limit) + 2.0: + # Still outside and getting FARTHER: halt again. Walking inbound + # (or sideways) is allowed -- that is how the robot comes home -- + # but each additional 2 m of escape re-trips the fence. + self._fence_halt_r = r + self._cancel_goto() + self._send({"action": "stop"}) + logger.warning( + f"[{self.robot_key}] GROUND FENCE: still receding ({r:.0f} m) -- halted again" + ) + + def _state_loop(self) -> None: + """Republish onto the shared fleet bus so the coordinator sees this robot.""" + period = 1.0 / STATE_PUBLISH_HZ + while self._running: + try: + self.drone_state.publish(String(json.dumps(self._state_dict()))) + except Exception as e: + logger.debug(f"[{self.robot_key}] state publish failed: {e}") + time.sleep(period) + + # -- state ------------------------------------------------------------- + + def _state_dict(self) -> dict[str, Any]: + with self._lock: + raw = dict(self._latest) + if not raw: + return { + "key": self.robot_key, + "robot_class": "legged", + "connected": False, + "ts": time.time(), + } + # Simulator world is NWU; the fleet stack reasons in NED. + x, y, z = raw.get("nwu", [0.0, 0.0, 0.0]) + vx, vy, vz = raw.get("velocity", [0.0, 0.0, 0.0]) + north, east = x, -y + lat, lon = offset_latlon(self.config.origin_lat, self.config.origin_lon, north, east) + return { + "key": self.robot_key, + "robot_class": "legged", + "connected": True, + "capabilities": ["ground", "camera"], + "ned": [north, east, -z], + "velocity_ned": [vx, -vy, -vz], + # Same shape the drones publish, so SwarmCoordinator's Haversine + # geometry covers ground robots without special-casing them. + "global": {"lat": lat, "lon": lon, "rel_alt_m": z}, + # Simulator yaw is about +z in NWU (North toward West); NED yaw runs + # North toward East, so it is the negation. + "yaw_ned_rad": -float(raw.get("yaw_rad", 0.0)), + "walking": raw.get("walking", False), + "fallen": raw.get("fallen", False), + "altitude_m": z, + "height_m": raw.get("height_m", z), + "nominal_height_m": raw.get("nominal_height_m", self.config.nominal_height_m), + # No battery model in sim yet; report the placeholder the PDF asks + # for rather than omitting the field and breaking fleet_state. + "battery_pct": 100.0, + "ts": time.time(), + } + + @skill + def state(self) -> str: + """Report this robot's position, height and connection state.""" + return json.dumps(self._state_dict(), indent=2) + + # -- skills ------------------------------------------------------------ + + @skill + def stand(self) -> str: + """Return to the nominal standing stance. + + If the robot is down this also attempts a recovery: the fall latch is + cleared and the legs driven to full stance, which rights the body when it + has toppled onto its side. A fully inverted robot cannot get up -- these + legs have no abduction and there is no getup routine -- and the latch + simply re-arms. + """ + self._cancel_goto() + if self._latest.get("fallen"): + self._send({"action": "recover"}) + return f"{self.robot_key}: down — attempting recovery to stance" + self._send({"action": "stance", "height": 1.0}) + return f"{self.robot_key}: standing" + + def _is_go1(self) -> bool: + """True when the bridge is running the real Go1 + trained policy.""" + return self._latest.get("model") == "unitree_go1" + + @skill + def crouch(self) -> str: + """Lower the body to a stable crouch. + + On the real Go1 this is refused honestly: the trained policy is a + velocity controller with no height input, so the bridge would accept + the command and do nothing -- a false success worse than a refusal. + """ + if self._is_go1(): + return ( + f"{self.robot_key}: crouch not supported by the Go1 policy " + "(no height control); the robot stays standing" + ) + self._send({"action": "stance", "height": 0.5}) + return f"{self.robot_key}: crouching" + + @skill + def walk( + self, speed_mps: float = 0.25, turn_rate_rads: float = 0.0, lateral_mps: float = 0.0 + ) -> str: + """Walk at a body velocity until told otherwise. + + Args: + speed_mps: forward speed. Positive is forward; reverse is capped at + roughly half of forward, because this gait topples if pushed + backwards at full stride. + turn_rate_rads: yaw rate in the NED convention shared with the + aircraft -- positive is nose-right. Turning while walking + is far more effective than turning on the spot -- 2 DOF per leg + and no hip abduction means yaw comes only from a differential + stride. + """ + self._cancel_goto() + self._send( + { + "action": "walk", + "vx": float(speed_mps), + "vy": float(lateral_mps), + "wz": float(turn_rate_rads), + } + ) + lat = f", strafe {lateral_mps:+.2f} m/s" if abs(lateral_mps) > 1e-6 else "" + note = ( + " (strafe needs the Go1 policy; the primitive gait ignores it)" + if abs(lateral_mps) > 1e-6 and not self._is_go1() + else "" + ) + return ( + f"{self.robot_key}: walking at {speed_mps:.2f} m/s, " + f"turn {turn_rate_rads:+.2f} rad/s (+ = right){lat}{note}" + ) + + @skill + def halt(self) -> str: + """Stop walking and settle into a standing stance. + + Named ``halt``, not ``stop``: ``Module.stop()`` is the framework's + teardown RPC (it closes the module's RPC, tools and event loop). A skill + called ``stop`` shadows it, so the module never tears down -- and worse, + anything that called ``self.stop()`` internally would silently kill the + module instead of halting the robot. + """ + self._cancel_goto() + self._send({"action": "stop"}) + return f"{self.robot_key}: stopped" + + def _move_target(self, forward_m: float, right_m: float) -> tuple[float, float, float] | None: + """(target_n, target_e, yaw_now) from the LIVE pose, or None if not connected.""" + state = self._state_dict() + if not state.get("connected"): + return None + n, e, _ = state["ned"] + yaw = float(state.get("yaw_ned_rad", 0.0)) + # NED, yaw from north toward east: heading = (cos y, sin y), + # right-hand vector = (-sin y, cos y). + tn = n + forward_m * math.cos(yaw) - right_m * math.sin(yaw) + te = e + forward_m * math.sin(yaw) + right_m * math.cos(yaw) + return tn, te, yaw + + @skill + def move(self, forward_m: float = 0.0, right_m: float = 0.0) -> str: + """Walk a body-relative offset from where the robot is RIGHT NOW, + WITHOUT changing which way it faces. + + THE tool for "move left 5 meters", "go forward 2", "back up a bit": + the target comes from the robot's live pose at call time, so + consecutive relative commands chain from wherever the previous one + ended -- and the heading is held (the Go1 policy strafes), so + "forward" after a "left" still means the direction the robot was + facing. ``goto`` is different on purpose: it turns toward its target + (right for long distances, but it changes what "forward" means). + + The control loop runs INSIDE the simulator at physics rate. A first + version ran here on wall-clock ticks and at ~40x realtime every + command acted on 6-12 sim-seconds of stale state -- the robot + wandered 140 m off a 5 m strafe. Control lives with the plant. + + Args: + forward_m: metres along the current heading (negative = back). + right_m: metres to the robot's right (negative = LEFT). + """ + tgt = self._move_target(forward_m, right_m) + if tgt is None: + return f"{self.robot_key}: NOT CONNECTED (no state from the bridge yet)" + tn, te, yaw0 = tgt + r = math.hypot(tn, te) + if r > self.config.max_range_m: + return ( + f"{self.robot_key}: REJECTED move -- target ({tn:.0f}, {te:.0f}) is " + f"{r:.0f} m from origin, beyond the {self.config.max_range_m:.0f} m radius" + ) + self._cancel_goto() + self._send({"action": "move_to", "north": tn, "east": te, "yaw": yaw0}) + return ( + f"{self.robot_key}: moving {forward_m:+.1f} m forward / {right_m:+.1f} m right " + f"(heading held) -> NED ({tn:.1f}, {te:.1f}); poll state until walking=false" + ) + + @skill + def goto(self, north: float = 0.0, east: float = 0.0) -> str: + """Walk to a point in the shared world frame, steering as it goes. + + With the real Go1 and its trained policy this is dependable: measured + 6/6 spread targets reached, worst arrival error 0.62 m, including + targets directly behind the robot. On the primitive fallback gait it + reaches about 4 of 6 and can topple on large sustained turns -- there + ``walk``/``halt`` remain the reliable primitives. + + Targets outside the fleet operating radius (``max_range_m``, shared + with the aircraft coordinator) are refused at dispatch. + + Runs closed-loop in the background and returns immediately; poll + ``state`` or ``fleet_state`` to watch progress. Calling ``goto``, + ``walk`` or ``stop`` again cancels an in-flight one. + """ + + r = math.hypot(north, east) + if r > self.config.max_range_m: + return ( + f"{self.robot_key}: REJECTED goto ({north:.0f}, {east:.0f}) -- " + f"{r:.0f} m from origin exceeds the {self.config.max_range_m:.0f} m " + "operating radius" + ) + if not self._latest: + return f"{self.robot_key}: no position yet from the simulator" + self._cancel_goto() + self._goto_stop.clear() + self._goto_thread = threading.Thread( + target=self._goto_loop, + args=(float(north), float(east)), + daemon=True, + name=f"legged-goto-{self.robot_key}", + ) + self._goto_thread.start() + return f"{self.robot_key}: walking to NED ({north:.1f}, {east:.1f})" + + def _cancel_goto(self) -> None: + if self._goto_thread is not None and self._goto_thread.is_alive(): + self._goto_stop.set() + self._goto_thread.join(timeout=1.0) + self._goto_thread = None + + def _goto_loop(self, north: float, east: float) -> None: + deadline = time.monotonic() + GOTO_TIMEOUT_S + while not self._goto_stop.is_set() and time.monotonic() < deadline: + state = self._state_dict() + if not state.get("connected"): + time.sleep(GOTO_TICK_S) + continue + if state.get("fallen"): + logger.warning(f"[{self.robot_key}] goto aborted: robot is down") + self._send({"action": "stop"}) + return + cur_n, cur_e, _ = state["ned"] + dn, de = north - cur_n, east - cur_e + distance = math.hypot(dn, de) + if distance <= GOTO_ARRIVE_M: + self._send({"action": "stop"}) + logger.info(f"[{self.robot_key}] goto arrived ({distance:.2f} m)") + return + # Heading error, wrapped to [-pi, pi] so it turns the short way. + bearing = math.atan2(de, dn) + yaw_err = math.atan2( + math.sin(bearing - state["yaw_ned_rad"]), + math.cos(bearing - state["yaw_ned_rad"]), + ) + # Hold cruise speed even when badly misaligned; the turn is only + # effective while moving. Ease off on the final approach. + if abs(yaw_err) < GOTO_YAW_DEADBAND_RAD: + yaw_err = 0.0 + speed = GOTO_CRUISE_MS + if distance < GOTO_SLOWDOWN_M: + speed = max(GOTO_MIN_SPEED_MS, GOTO_CRUISE_MS * distance / GOTO_SLOWDOWN_M) + self._send({"action": "walk", "vx": speed, "wz": GOTO_YAW_GAIN * yaw_err}) + time.sleep(GOTO_TICK_S) + self._send({"action": "stop"}) + if not self._goto_stop.is_set(): + logger.warning(f"[{self.robot_key}] goto timed out after {GOTO_TIMEOUT_S:.0f}s") + + @skill + def set_height(self, height_fraction: float = 1.0) -> str: + """Set stance height as a fraction of nominal (0.35 = lowest, 1.0 = standing).""" + if self._is_go1(): + return ( + f"{self.robot_key}: set_height not supported by the Go1 policy " + "(no height control); the robot stays at its trained stance" + ) + h = max(0.35, min(1.0, float(height_fraction))) + self._send({"action": "stance", "height": h}) + return f"{self.robot_key}: stance height -> {h:.2f} of nominal" + + # -- fleet bus --------------------------------------------------------- + + def _on_swarm_cmd(self, msg: String) -> None: + """Act on coordinator broadcasts addressed to this robot (or to `all`). + + Shared verbs only. `land`/`rtl` mean "get low and stay put" for a ground + robot, which is the closest honest equivalent of what they mean for an + aircraft; motion verbs that have no ground analogue are ignored rather + than faked. + """ + try: + cmd = json.loads(msg.data if hasattr(msg, "data") else str(msg)) + except (ValueError, AttributeError): + return + target = cmd.get("target", "all") + if target not in ("all", self.robot_key): + return + action = cmd.get("action", "") + if action in ("land", "emergency_land", "hold"): + logger.info(f"[{self.robot_key}] swarm_cmd {action}: stopping and crouching") + self.halt() + self.crouch() + elif action == "ground_goto": + # The coordinator's staged missions route ground robots with this. + # It reuses the public goto, so the operating radius and fall + # handling apply exactly as if the operator had called it. + # Coordinates live under "args" -- _send() nests its kwargs there. + # Reading them from the top level silently yielded the (0, 0) + # defaults, so the dog dutifully walked home instead of to its + # waypoint and the mission hung waiting for an arrival that could + # never happen. + args = cmd.get("args") or {} + self.goto(north=float(args.get("north", 0.0)), east=float(args.get("east", 0.0))) + elif action in ("takeoff", "rtl"): + logger.info(f"[{self.robot_key}] swarm_cmd {action}: standing") + self.stand() + elif action == "kill": + logger.warning(f"[{self.robot_key}] swarm_cmd kill: collapsing to the ground") + self._send({"action": "stance", "height": 0.35}) + + +__all__ = ["STATE_PUBLISH_HZ", "LeggedSimConfig", "LeggedSimModule"] From 5f4ee149acebfc066c18e03d4e71972ad90f6e2f Mon Sep 17 00:00:00 2001 From: Ezenbaatar Batjargal Date: Mon, 31 Aug 2026 11:38:36 -0700 Subject: [PATCH 4/4] feat(mcp): per-instance tool naming, OpenAI-safe aliases, resilient agent loop Skills from namespaced module instances qualify as instance/module/skill only when ambiguous; the LLM client shows sanitized aliases (OpenAI rejects '/' in function names, and one bad name 400s the whole request) while calling the server by the real name; the agent message loop survives a failed turn instead of dying silently behind a 'thinking...' UI. Co-Authored-By: Claude Fable 5 --- dimos/agents/mcp/mcp_client.py | 44 ++++++++++++++++-------- dimos/agents/mcp/mcp_server.py | 52 +++++++++++++++++++++++++---- dimos/agents/mcp/test_mcp_server.py | 31 +++++++++++++++++ dimos/core/module.py | 8 ++++- 4 files changed, 113 insertions(+), 22 deletions(-) diff --git a/dimos/agents/mcp/mcp_client.py b/dimos/agents/mcp/mcp_client.py index e0394c3217..f7f91352bb 100644 --- a/dimos/agents/mcp/mcp_client.py +++ b/dimos/agents/mcp/mcp_client.py @@ -14,22 +14,15 @@ from collections.abc import Callable from queue import Empty, Queue +import re from threading import Event, RLock, Thread import time from typing import Any import uuid -import warnings - -from langchain_core._api.deprecation import LangChainPendingDeprecationWarning - -# Importing langchain_core un-mutes its pending-deprecation warnings, so this ignore -# must be registered after that import to take precedence. It silences the noisy -# `allowed_objects` warning emitted when langchain.agents pulls in langgraph.checkpoint. -warnings.filterwarnings("ignore", category=LangChainPendingDeprecationWarning) from langchain.agents import create_agent from langchain.chat_models import init_chat_model -from langchain_core.messages import HumanMessage +from langchain_core.messages import AIMessage, HumanMessage from langchain_core.messages.base import BaseMessage from langchain_core.tools import StructuredTool from langchain_openai import ChatOpenAI @@ -189,7 +182,16 @@ def _try_fetch_tools(self, timeout: float, interval: float) -> dict[str, Any] | def _mcp_tool_to_langchain(self, mcp_tool: dict[str, Any]) -> StructuredTool: name = mcp_tool["name"] + # OpenAI restricts function names to [a-zA-Z0-9_-]. Namespaced MCP + # tools ("dog1/leggedsimmodule/state") contain '/', and ONE such name + # 400-fails the entire request -- the model never sees any tool and, + # before the guard in _thread_loop, the agent thread died on it. The + # model is shown a sanitized alias; the MCP server is always called + # with the real name via the closure below. + model_name = re.sub(r"[^a-zA-Z0-9_-]", "_", name) description = mcp_tool.get("description", "") + if model_name != name: + description = f"[{name}] {description}" input_schema = mcp_tool.get("inputSchema", {"type": "object", "properties": {}}) def call_tool(**kwargs: Any) -> str: @@ -210,7 +212,7 @@ def call_tool(**kwargs: Any) -> str: return text return StructuredTool( - name=name, + name=model_name, description=description, func=call_tool, args_schema=input_schema, @@ -335,10 +337,24 @@ def _thread_loop(self) -> None: except Empty: continue - with self._lock: - if not self._state_graph: - raise ValueError("No state graph initialized") - self._process_message(self._state_graph, message) + try: + with self._lock: + if not self._state_graph: + raise ValueError("No state graph initialized") + self._process_message(self._state_graph, message) + except Exception as e: + # A single failed turn (network blip, provider 4xx/5xx) must + # not kill this thread: before this guard, one exception left + # every later message queued forever while the UI showed + # "thinking...". Tell the operator and keep serving. + logger.error(f"Agent turn failed: {type(e).__name__}: {e}") + try: + self.agent.publish( + AIMessage(content=f"[agent error] {type(e).__name__}: {e}") + ) + self.agent_idle.publish(True) + except Exception: + pass def _process_message( self, state_graph: CompiledStateGraph[Any, Any, Any, Any], message: BaseMessage diff --git a/dimos/agents/mcp/mcp_server.py b/dimos/agents/mcp/mcp_server.py index 61d7572af8..1ea862a9c5 100644 --- a/dimos/agents/mcp/mcp_server.py +++ b/dimos/agents/mcp/mcp_server.py @@ -91,14 +91,49 @@ def _handle_initialize(req_id: Any) -> dict[str, Any]: ) +def tool_names(skills: list[SkillInfo]) -> dict[str, SkillInfo]: + """Map MCP tool name -> skill, disambiguating duplicates by instance. + + A skill name is left bare when only one module offers it, which is the + common case and keeps existing blueprints and docs working. When several + module *instances* offer the same skill -- a namespaced fleet of three + drones all exposing ``takeoff`` -- the bare name is ambiguous and the last + registration would silently win, so those become + ``drone1/px4dronemodule/takeoff``. That matches the RPC topic exactly. + """ + counts: dict[str, int] = {} + for s in skills: + counts[s.func_name] = counts.get(s.func_name, 0) + 1 + out: dict[str, SkillInfo] = {} + for s in skills: + name = s.func_name if counts[s.func_name] == 1 else f"{s.class_name}/{s.func_name}" + out[name] = s + return out + + +def with_qualified_aliases(named: dict[str, SkillInfo]) -> dict[str, SkillInfo]: + """Extend a tool-name map with each skill's instance-qualified name. + + ``tools/list`` shows only the canonical names from `tool_names`, but a call + to ``dog1/go2firstcontact/safe_move`` must work even while that dog is the + only module offering ``safe_move``: the qualified form is the RPC topic, so + it is never ambiguous, and an operator's fleet scripts keep working when + the fleet is temporarily one robot. + """ + out = dict(named) + for s in named.values(): + out.setdefault(f"{s.class_name}/{s.func_name}", s) + return out + + def _handle_tools_list(req_id: Any, skills: list[SkillInfo]) -> dict[str, Any]: tools = [] - for s in skills: + for name, s in tool_names(skills).items(): schema = json.loads(s.args_schema) description = schema.pop("description", None) schema.pop("title", None) - tool: dict[str, Any] = {"name": s.func_name, "inputSchema": schema} + tool: dict[str, Any] = {"name": name, "inputSchema": schema} if description: tool["description"] = description if s.uses or s.lifecycle != "instant": @@ -384,12 +419,15 @@ def on_system_modules(self, modules: list[RPCClient]) -> None: app.state.skills = [ skill_info for module in modules for skill_info in (module.get_skills() or []) ] - app.state.skills_by_name = {s.func_name: s for s in app.state.skills} + # Keyed by MCP tool name, which is instance-qualified when a skill + # exists on more than one module instance. The RpcCall still targets the + # bare func_name on that instance's RPC topic. Qualified aliases are + # callable but not listed (see with_qualified_aliases). + named = with_qualified_aliases(tool_names(app.state.skills)) + app.state.skills_by_name = named app.state.rpc_calls = { - skill_info.func_name: RpcCall( - None, self.rpc, skill_info.func_name, skill_info.class_name, [] - ) - for skill_info in app.state.skills + name: RpcCall(None, self.rpc, skill_info.func_name, skill_info.class_name, []) + for name, skill_info in named.items() } @skill diff --git a/dimos/agents/mcp/test_mcp_server.py b/dimos/agents/mcp/test_mcp_server.py index 0e2a74925b..a5da0cd889 100644 --- a/dimos/agents/mcp/test_mcp_server.py +++ b/dimos/agents/mcp/test_mcp_server.py @@ -334,3 +334,34 @@ def test_instant_holder_conflict_waits_then_runs() -> None: app.state.skills_by_name = saved_skills app.state.cap_registry = saved_registry app.state.cap_acquire_timeout = saved_timeout + + +def test_tool_names_qualified_aliases() -> None: + """A unique skill lists bare but stays callable by its qualified RPC-topic + name, so fleet scripts (dog1/go2firstcontact/safe_move) keep working when + only one robot offers the skill. Duplicates stay qualified-only.""" + from dimos.agents.mcp.mcp_server import tool_names, with_qualified_aliases + + schema = json.dumps({"type": "object", "properties": {}}) + skills = [ + SkillInfo(class_name="dog1/go2firstcontact", func_name="safe_move", args_schema=schema), + SkillInfo(class_name="dog1/go2connection", func_name="stand_up", args_schema=schema), + ] + + named = tool_names(skills) + assert set(named) == {"safe_move", "stand_up"} # unique skills list bare + + callable_names = with_qualified_aliases(named) + assert callable_names["safe_move"] is callable_names["dog1/go2firstcontact/safe_move"] + assert callable_names["stand_up"] is callable_names["dog1/go2connection/stand_up"] + + # Two instances offering the same skill: bare name is ambiguous and gone, + # aliasing adds nothing beyond the already-qualified canonical names. + duplicated = skills + [ + SkillInfo(class_name="dog2/go2firstcontact", func_name="safe_move", args_schema=schema) + ] + named_dup = tool_names(duplicated) + assert "safe_move" not in named_dup + assert set(with_qualified_aliases(named_dup)) == set(named_dup) | { + "dog1/go2connection/stand_up" + } diff --git a/dimos/core/module.py b/dimos/core/module.py index f5aab99421..81e827b7df 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -476,7 +476,13 @@ def get_skills(self) -> list[SkillInfo]: lifecycle = getattr(attr, "__skill_lifecycle__", "instant") skills.append( SkillInfo( - class_name=self.__class__.__name__, + # The instance name when this module is namespaced, so a + # fleet of identical modules stays individually + # addressable. RpcCall uses this as the RPC topic + # prefix, and _rpc_name() picks the same string, so the + # two agree. Falling back to the class name keeps + # single-instance blueprints unchanged. + class_name=self.config.instance_name or self.__class__.__name__, func_name=name, args_schema=schema, uses=uses,