FlyBy: GPU-Accelerated Reinforcement Learning for Quadrotors

This is a web write-up of my final-year project in the Department of Aeronautics at Imperial College London (2024/2025): High-Performance Computing for Aerial Autonomy: GPU-Driven Training of Multirotor Vehicles. The full paper is linked at the end.

The short version: most reinforcement-learning libraries are built for ease of use, not speed. For a drone, that trade-off hurts — training a single control policy can take hours, and anything that needs many trained agents (like co-design, where you evolve a vehicle's shape and its controller together) becomes computationally hopeless. FlyBy is my attempt to fix that: a quadrotor RL framework written end-to-end in JAX so that the simulator, the policy, and the training loop all run on the GPU at once. The result trains policies in minutes instead of hours, while using a more physically honest model of how rotors actually produce thrust.

The FlyBy loop. A vectorised JAX simulator runs up to 160 environments in parallel; each emits a 16-dimensional observation, the actor maps it to a 4-dimensional action (collective thrust and three body torques), and the action feeds back into the simulator. The critic and curriculum manager exist only during training.

Why speed is the whole game

The motivation comes from co-design — evolving an agent's form and function at the same time, the way natural selection does. You train a population of agents, keep the best, mutate them, and train again. Because a good shape depends on a good controller and vice versa, there's no shortcut: it's a brute-force, multi-objective search that needs an enormous number of training runs.

That only works if a single training run is cheap. And on conventional stacks it isn't — training time grows with the size of the state space, the complexity of the objective, and the range of conditions the agent has to generalise across. Non-linear flight dynamics make it worse, because behaviour is path-dependent. So the bottleneck isn't the algorithm; it's the sheer throughput of the simulate-act-learn loop. FlyBy's goal was simple to state: be faster than any current open-source option.

The framework

FlyBy has four parts: a simulation environment (a high-fidelity quadrotor model), a policy network, a training engine (PPO), and a curriculum manager that ramps up task difficulty automatically.

The environment is implemented against the Gymnax interface so it stays JAX-compatible:

• Observation — 16 dimensions: position, linear velocity, orientation (Euler angles), angular velocity, the normalised direction to the next waypoint, and the current waypoint index. (The underlying rigid-body state is 12-dimensional; the extra four numbers are what the agent is allowed to see.)
• Action — 4 dimensions: collective thrust plus the three body torques, [T, τ_x, τ_y, τ_z]. Thrust is normalised to [0, 1] and scaled so that T_max = 2·m·g; torques are normalised to [-1, 1].
• Physics: full 6-DOF rigid-body dynamics with semi-implicit Euler integration at 50 Hz (Δt = 0.02 s).

One deliberate design choice: FlyBy applies thrust and torque to the body directly, skipping the usual control-allocation "mixer" that converts those four commands into four individual motor speeds. Dropping the mixer raises the level of abstraction, makes learning easier, and lets the same policy generalise across different quadrotor layouts. The continuous-time dynamics it integrates are the standard ones:

The policy and value networks are an actor-critic pair built in Flax — four hidden layers of 160 units each, tanh activations, orthogonal initialisation, shared feature-extraction layers, and a diagonal Gaussian policy with learnable log standard deviations. At inference, only the actor flies the drone; the critic is purely a training aid.

Physically grounded aerodynamics

Almost every quadrotor RL paper models a rotor with the quadratic approximation T = k_T·ω² — thrust proportional to the square of rotor speed. It's convenient and wrong in interesting ways, especially in forward flight.

FlyBy instead uses Blade Element Momentum (BEM) theory. The blade is sliced into annular elements; axial and tangential induction factors are solved by fixed-point iteration, with Prandtl tip- and hub-loss corrections and a Buhl/Burton empirical correction in the high-induction regime. That iteration is far too expensive to run inside the training loop, so FlyBy precomputes converged thrust and torque over a grid of axial velocity and rotor speed, then looks them up at runtime with bilinear interpolation.

The payoff: the lookup table reproduces the full BEM solver to within R² over 0.99, runs 150–300× faster than the full calculation, and keeps prediction accuracy above 99%. Validated against static-thrust measurements for the APC 10×4.7 propeller, the BEM curve tracks the experimental data closely across the operating range:

BEM-predicted thrust (blue) against experimental measurements (red) across rotor speed. Close agreement, at a fraction of the cost of the full solver.

Putting everything on the GPU

This is where the speed comes from. The entire pipeline is built on JAX's four core transformations — jit to compile for the hardware, vmap to vectorise across environments, grad for gradients, and scan for sequences. Because the simulator is written functionally (indexed updates, branch-free control flow with jnp.where and jax.lax.cond, splittable RNG keys), over 99% of the environment step compiles straight to the GPU with almost no host-device chatter.

The headline trick is running many environments at once. A single vmap turns one environment into a batch of them:

obsv, env_state, reward, done, info = jax.vmap(
    env.step, in_axes=(0, 0, 0, None)
)(rng_step, env_state, action, env_params)

With up to 160 parallel environments, FlyBy reaches over 150,000 environment-steps per second on consumer hardware (an NVIDIA Tesla T4), against roughly 1,000 steps per second for the same logic running serially on CPU. The per-step GPU/CPU speedup scales almost linearly with the number of parallel environments before plateauing as the GPU's cores and memory saturate:

Per-step speedup (GPU vs CPU) as a function of parallel environment count — from a few× at low counts up to roughly 110× at 160 environments, then flattening.

Profiling the loop shows why vectorising the environment matters so much: about 45% of per-step time is environment simulation, 35% is neural-network work, and 20% is overhead. The policy is optimised with PPO — clipping ε = 0.2, discount γ = 0.99, GAE λ = 0.95, learning rate 8e-4 with linear annealing, four epochs and four minibatches per update — maximising the standard clipped surrogate objective:

End to end, that adds up to up to 11× faster training than the CPU baseline on a hover task, and on a nine-gate racing track FlyBy with CUDA trained in about four minutes what Stable-Baselines3 needed roughly 18 hours for — a 270× reduction.

Learning to race with a curriculum

Throughput buys you nothing if the agent never learns the hard task. Asked to fly a full figure-8 from scratch, training simply fails — the reward is too sparse and the dynamics too unforgiving for the policy to ever stumble onto success.

The fix is curriculum learning: eight difficulty levels (labelled 0, 1, 1.3, 1.5, 1.7, 2, 3, and 4) that start with a single, generously-thresholded waypoint and build up to a tight, twelve-waypoint figure-8. The curriculum manager promotes the agent to the next level automatically once its success rate over a recent window clears a level-specific bar — and those bars loosen as the task gets harder (from 0.8 at the easiest level down to 0.4), reflecting the rising difficulty.

The contrast is stark:

Left: with a curriculum, success rate climbs and the difficulty level steps up automatically (0 → 1 → 1.3 → 1.5 → 1.7 …). Right: trained directly at difficulty 3, the same agent never gets off the floor.

With the curriculum, the agent reaches a success rate above 0.8 at the highest difficulty; trained directly at maximum difficulty, it stays near zero throughout. The trained policy then flies the full track — threading every gate in sequence:

A learned policy completing the difficulty-4 figure-8 gate track. Reported success rates were 82% on the figure-8 and 73% on a harder nine-gate track.

Results at a glance

A few of the ablations that shaped the final design:

Knob	Finding
Learning rate	Annealed schedules reached episode rewards of 1,500–2,000, versus about 1,000 for a fixed rate.
Entropy coefficient	An inverted-U: about 0.01 was optimal; 0.1 over-explored and never cleared a reward of 500.
Episode length	700-step episodes peaked near 2,200 reward; 200-step episodes hit a ceiling around 650.
Target distance	A 2 m target reached about 4,500 reward; 5 m about 2,000; 10 m stayed below 500.
BEM lookup	150–300× faster than full BEM, R² over 0.99, accuracy above 99%.

What's next

FlyBy is quadrotor-only for now, the curriculum stages are hand-designed, and the aerodynamic model still ignores rotor-rotor interactions, ground effect, and turbulence. The biggest open item is the obvious one: none of this has been validated on real hardware yet, so the sim-to-real story is still a hypothesis. Natural next steps are automated curriculum generation, multi-GPU training, and combining the framework with imitation or model-based methods for even better sample efficiency. The broader bet is that making each training run cheap is what finally makes co-design — evolving the drone and its brain together — practical.

Read the full paper

The complete write-up, with all derivations, tables, and references, is here: FlyBy — full paper (PDF).