Aerial manipulation: keep the arm pretrained, learn flight separately

Problem

• Aerial data is scarce. Tabletop manipulation data can be collected at scale. Aerial manipulation data needs expert teleoperation or hand-coded scripts and MPC, which is not enough to train a whole-vehicle policy from scratch.
• The arm is already solved elsewhere. A VLA pretrained on tabletop data already has the language understanding and fine arm control. Scarce flight data should not be spent retraining it.
• Don't edit the arm's actions with flight constraints. Writing flight feasibility into one shared action chunk changes the pretrained arm behaviour, and ties the flight coupling to that one backbone, so upgrading the VLA later means rebuilding the constraint.

Idea

• Two policies. A small Flight DiT handles flight. A mature VLA (π0.5) handles the arm. Flight is learned by the small model; the VLA only learns to take the flight into account.
• Flight DiT. A lightweight diffusion Transformer, sized like RDT-170M to train on one 4090. It reads language, images and proprioception, and outputs a flight chunk a_flight = [Δx, Δy, Δz, Δψ] plus a flight-intent token F_flight.
• Arm VLA. F_flight is encoded by a small projector and injected into π0.5, at an intermediate transformer layer or at the action expert, so the flight intent is present at every denoising step. π0.5 is then fine-tuned with all of its parameters (not LoRA) on the aerial demonstrations with the token in place, and outputs a_arm = [q1, q2, q3, q4, gripper] (FA-Hexa's arm: three pitch joints and a wrist roll).
• Easy to swap. The two policies meet only at the injected token. Replacing the arm VLA means re-training the projector and the fine-tune on the new backbone; the Flight DiT and its data stay as they are.
• Low level stays split. The flight controller tracks the flight chunk; joint servos track the arm.

Pipeline

• Stage A: flight only. Train the Flight DiT on a_flight and extract F_flight.
• Stage B: inject and fine-tune. Encode F_flight, inject it into π0.5 (intermediate layer or action expert) and fine-tune all of π0.5 on the same demonstrations.
• Stage C: deploy. Run both chunks in sync; flight and arm are tracked by their own low-level controllers.

Related work

Models

• Transformer (Vaswani et al., NeurIPS 2017): the backbone of both policies. Cross-attention is where language, images and F_flight enter.
• DiT (Peebles & Xie, ICCV 2023): a Transformer inside diffusion. The Flight DiT uses the same idea to generate a short flight-action chunk.
• RDT-1B / RDT-170M (Liu et al., ICLR 2025): a diffusion Transformer robot policy. RDT-1B fine-tuning costs about a week on 4×A100; RDT-170M (hidden 1024, depth 14, 64-step prediction) fits on one 4090 and is the starting point for the Flight DiT.

Aerial manipulation

• Flying Hand (He et al., 2025): an end-effector-centric interface on a fully actuated hexarotor with whole-body MPC. The high level only talks about the end effector. FA-Hexa is the same class of vehicle.
• Aerial Manipulation in the Wild (Zhan et al., 2026): a Diffusion Policy proposes end-effector motion and whole-body MPC executes it, outdoors. Here the upper model is a tabletop-pretrained VLA instead.
• π, But Make It Fly / AirVLA (Tucker et al., 2026): visual features transfer to flight, control dynamics do not. Their fix writes payload guidance into the VLA's sampling (23% → 50% pick-and-place). This project keeps the lesson but not the mechanism: flight is a separate module and only a condition.
• AC-DiT (Chen et al., 2025): extract base motion first, then condition the body on it. F_flight comes from this idea, but AC-DiT still predicts the whole body with one network (about two weeks on 8×A100).

Plan

• Goal: ICML 2027. Dates are not out yet; going by the 2026 cycle, the abstract is due around 23 Jan and the full paper around 28 Jan 2027 (2026 dates). That is about 15 weeks from mid-October.
• Done: the dataset. 4,800 verified demonstrations, 12 tasks × 5 conditions × 80, recorded with the upgraded specialists (Parts 2 and 3).
• Next: the Flight DiT (Part 4). Stage A on the 120 Hz flight data, then Stage B: inject F_flight into π0.5 and fine-tune it fully.
• Then: the benchmark (Part 5). Diffusion Policy, π0.5 fine-tuned with LoRA, and the Flight DiT + fully fine-tuned π0.5, all trained on the same demonstrations and evaluated with AM-Bench's success checks.

Specialist dataset on FA-Hexa

• “Official” means AM-Bench. The unmodified tasks, scripted specialists and success checks of AM-Bench (paper, code). “Upgraded” is the same code with the changes below.
• Why check the specialists. π0.5 is trained on their successful demonstrations, so a specialist that cheats a loose check or fails for unrelated reasons would go straight into the data. Every task was recorded and watched before recording at scale.
• Setup. FA-Hexa (fully actuated hexarotor, 4-DoF arm), absolute end-effector targets, IK plus PID at 120 Hz. Clips are 1280×720, real time, and stop 5 s after success.

Notes

• Disturbances. AM-Bench offers thrust saturation, ground and near-wall effects, drag, a constant wind force and sensor/actuator noise, all off by default; no gusts, wind torque or payload changes.
• π0.5 data. 80 demonstrations per task at 20 Hz in LeRobot format (EE camera, 8- or 12-value state, 50-step action chunks); AM-Bench fine-tunes all parameters, which needs more than 70 GB of GPU memory (LoRA: more than 22.5 GB).
• AC-DiT data. RDT-170M mobility head plus RDT-1B whole-body head, trained on ManiSkill-HAB HDF5 at 20 Hz (three cameras, point cloud, 128-value state vector). An aerial base needs new state slots.

Full notes with file paths and sources: docs/research/disturbances-pi05-acdit-data.md.

Related work

• AM-Bench (Wang et al., 2026): a modular simulation suite and benchmark for aerial manipulation policy learning, built on Isaac Lab. It provides the vehicles (FA-Hexa among them), the 12 tasks and the scripted specialists used here. Project page · Documentation · Code.
• Isaac Lab (NVIDIA): the GPU simulation framework AM-Bench runs on (Isaac Sim 5.1).

Demonstration dataset

We record 80 successful demonstrations for each of the 12 AM-Bench tasks under each of the 5 conditions of Part 2 (AM-Bench default; thrust limits with aerodynamic effects; the same with a wind of 5, 10 or 20 % of the vehicle's weight), using the upgraded scripted specialists on FA-Hexa. Every episode is verified individually before it enters the training set. The set is exported for π0.5 fine-tuning in the OpenPI format and kept, with a 120 Hz low-dimensional companion, for the Flight DiT. All data was recorded in one day on a single RTX 4080 SUPER.

• Specialists. The Part 2 specialists, one per task, run closed loop on the measured object state (peg tip, handle, lemon, ball) and output an absolute end-effector target and a gripper command. Two Wind 20 % variants were added for the recording noise (PushSlider pushes the knob with the side of the closed gripper; WipeWindow aligns before a gentle press); RotateValve gets 30 s episodes under the disturbance conditions.
• Simulation. Isaac Lab at 120 Hz with inverse kinematics and PID control; cameras render every sixth step, so a 20 Hz image is the exact render of the state it is paired with. Four environments per process, 20 m apart, two processes per GPU.
• Recording. LeRobot v3.0, successful episodes only: a 20 Hz set with both 384 × 384 cameras (AV1 video), the 35-D state, the label and the executed action, and a 120 Hz set with state and labels at every control step. The task prompt is stored with every frame.
• Verification and selection. Each episode is checked (below); the first 80 valid episodes per pair are selected and the rest excluded. The π0.5 export (LeRobot v2.1, 20 Hz, lossless PNG) is itself checked frame by frame against the selected source episodes.

Noise injection

Following DART, the executed target is perturbed and the specialist's unperturbed target is kept as the label, so the demonstrations contain small deviations and the corrections back to the nominal path. Three noise processes with the same standard deviation (5 mm, 0.025 rad; half of AM-Bench's defaults) behave very differently on a hovering vehicle (Figure 2). AM-Bench's per-step noise shakes the tool without moving the robot off its path, and every toss failed; an Ornstein–Uhlenbeck drift still changes the target every step, which the PID's derivative term turns into thrust spikes, and in 20 % wind the rotors saturate (press button 0 / 8). The smooth drift used here draws a new offset every 0.5 s and joins consecutive offsets with a smootherstep, so the target never jumps.

Recording settings

Verification

Every saved episode has to pass all of the following, and the two LeRobot sets are compared step by step. Rejected episodes stay on disk but are excluded from the selection and replaced by new recordings.

• Success. The task's success check fired and every sub-criterion was true at termination (a lemon that reached the bin without a registered grasp is rejected).
• Integrity. The 120 Hz episode has exactly the logged number of steps, the 20 Hz episode one frame per grid step, every video segment holds exactly its frames (LeRobot's streaming encoder silently drops a frame when its queue is full; the queue was enlarged after this was found), and the first, middle and last frames decode and differ.
• Clean start. The robot does not move by itself in the first 0.2 s. OpenDoor's door is sometimes spawned inside the robot's reset footprint and knocks the arm at its velocity limit; about a quarter of its episodes were rejected for this.
• Length. Within 0.6–1.5× the median of its task–condition pair, which excludes late or accidental successes.
• Metadata. The prompt, the disturbance switches and the wind force match the condition; the noise spread matches the setting.

Specialist success under the recording noise

Successes over attempts per task and condition, all recordings included. The median episode lasts 8.8 s in the default condition and 11.2 s in 20 % wind.

Table 1. Success rate in percent. Only successful episodes are stored; a low rate costs recording time, not data quality.

Data format

Table 2. Features of the 20 Hz set; the 120 Hz set holds the same episodes without images. The export keeps ee_image, base_image, ee_pos, ee_quat, gripper_width, actions and task, the schema of OpenPI's AM-Bench configuration, where actions become end-effector-relative 50-step chunks.

Code

Everything that produced the results lives in the private repository Aerial_Manipulation_Code; this page only shows the results.

• pipeline/datagen/: the batch recorder with GPU slot locks and re-recording, the per-episode verifier and selection, the OpenPI export and its check, and the one-command finish script.
• ambench/DATAGEN_CHANGES.diff: the recorder and exporter changes on top of AM-Bench (video storage, 20 Hz images with a 120 Hz companion, clean labels, smooth noise, one-pass export).
• docs/training-data.md: layout, features, rules and the full numbers. The data itself (303 GB) is not in the repository.

Related work

• π0.5: a Vision-Language-Action Model with Open-World Generalization (Physical Intelligence, 2025): the VLA fine-tuned on this data. OpenPI provides the training code and the AM-Bench data configuration used by the export (LeRobot v2.1, 20 Hz, 50-step action chunks).
• π0: A Vision-Language-Action Flow Model for General Robot Control (Black et al., 2024): the flow-matching action expert that π0.5 builds on.
• DART: Noise Injection for Robust Imitation Learning (Laskey et al., CoRL 2017): noise is injected into the supervisor's executed actions while the unperturbed action is recorded as the label, so the data covers the states a learner will drift into; the scheme used here.
• LeRobot (Cadene et al., 2024): the dataset format (v3.0 for recording, v2.1 for OpenPI) and the video encoding.

Flight DiT training

Coming soon.

Benchmark

Coming soon.