A recipe for drone racing with reinforcement learning | Home
This post is the third and final post in our series on quadcopter simulation.<br>The first post and second post derived and simulated the quadcopter’s equations of motion, first in 2D and then in 3D.<br>In this post we train a reinforcement learning policy to fly the quadcopter, first to hover at a fixed point, then to fly through a sequence of gates.
The previous posts were tutorial-like, we started from a free body diagram and, step-by-step, arrived to simulation code.<br>This post instead is more recipe-like, a collection of tricks and techniques I found useful to train an RL policy for quadcopter racing.<br>There is no rigorous proof of why these methods work, only empirical.
Because it is a recipe and not a walkthrough, the post shows only the code that matters for each decision.<br>The full environment is on GitHub at mrandri19/quadcopter-racing.
Here is the trained policy flying the eight-gate loop track.<br>The visualizations throughout the post are made with Rerun.
An incremental approach
Trying to implement a drone racing simulation, RL environment, reward, and model all at once just doesn’t work.<br>Too many things can be slightly off, affecting the policy’s performance, with no good way of debugging it.<br>So we split the task in three stages, each building on top of each other:
The “Hello, world! of RL”: training a policy to solve the inverted pendulum control problem with PPO and a vectorized stable-baselines3 env.
Building a custom quadcopter RL environment using MuJoCo, choosing action and observation spaces, splitting high-level RL-based and low-level P-controller, and a hovering reward.
Extending the environment from hovering to racing, with a new reward, observation, and random initialization design.
Solving inverted pendulum with vectorized PPO
Let’s start from the basics and setup our vectorized RL problem.<br>We use stable-baselines3’s PPO together with a vectorized InvertedPendulum-v5 (cartpole) environment, and check that the policy reaches the maximum episode length.
from stable_baselines3 import PPO<br>from stable_baselines3.common.env_util import make_vec_env<br>from stable_baselines3.common.policies import ActorCriticPolicy
def main() -> None:<br>env = make_vec_env("InvertedPendulum-v5", n_envs=128)<br>model = PPO(<br>policy=ActorCriticPolicy,<br>env=env,<br>learning_rate=3e-3, # 3e-4 (default lr) * sqrt(n_envs) ~= 3e-3<br>n_steps=512, # lower than default 2048, no need to have that many steps with 128 envs.<br>batch_size=1024, # higher than default 64 for better efficiency.<br>n_epochs=5, # lower than default 10, speeds up training, no need to refit 10 times<br>verbose=1,<br>model.learn(total_timesteps=750_000)
Let’s briefly discuss our choice of hyperparameters:
Rollout size is n_envs * n_steps: with 128 environments and 512 steps each, PPO collects 65,536 transitions before every policy update.<br>Making this batch large enough matters, as we want variety in the dataset we train on at each PPO update.
Number of PPO updates is total_timesteps / rollout_size, so roughly 11 updates here.<br>Given a large enough rollout size, the model also needs enough updates to “evolve” during training.<br>This is not supervised learning, where the dataset is fixed.<br>As the policy evolves, new patterns appear in the rollout data, so we want enough PPO updates to explore this space.
Number of gradient steps is total_timesteps * n_epochs / batch_size: each rollout is split into minibatches of batch_size and reused for n_epochs passes.<br>The bigger the batch size, the less noisy the gradient but the fewer training iterations we do per PPO update.
Results
On this environment, each alive step gives +1 reward, so episode length and reward will match exactly.<br>Training runs for 750k timesteps across 128 parallel environments, or roughly 11 PPO updates.<br>Episode length grows smoothly from the first update, and the policy is near the 1,000-step cap by 600k steps.
Recipe: validate PPO and its hyperparameters on a toy task first, because it separates “the algorithm is misconfigured” from “my environment is wrong”.
Learning to hover with a custom quadcopter environment
Now that we have shown that PPO plus our hyperparameters works on a toy task, let’s develop the quadcopter environment and setup a simple hover reward.<br>We will use MuJoCo as our physics simulator.
As said in the introduction, we won’t go through all the code necessary to implement the environment, only the parts worth discussing.<br>The rest is in src/quadcopter_racing/part_2.py.
Action design: CTBR + P controller
The policy does not directly output four motor commands.<br>Instead, it outputs a collective thrust and body rates (CTBR) command: mass-normalized thrust and desired roll, pitch, and yaw rates.<br>This is the action design that several papers such as Champion-level drone racing using deep reinforcement learning or Deep Drone Acrobats use.<br>A low-level proportional (P) controller then converts this into individual motor commands,...