Reinforcement Learning for Drone Control
Reinforcement learning teaches drones to fly through trial and error, much like how a bird learns to navigate by exploring its environment. Instead of hand-coding control rules, the drone discovers optimal behaviors by maximizing cumulative rewards.
The RL Agent-Environment Loop
Sponsored
RL operates on a simple principle: an agent takes actions in an environment, receives rewards, and learns to maximize long-term reward.
RL Agent-Environment Interaction
Q-Learning for Drones
Q-learning is a value-based method that learns the quality (Q-value) of state-action pairs. For drones, discretized states and actions make Q-learning applicable to simple tasks.
Q-Learning Update Rule
Architecture Diagram
Q(s, a) ← Q(s, a) + α[r + γ·max Q(s', a') - Q(s, a)]
python
from collections import defaultdict
class QLearningDrone:
def __init__(self, state_bins=20, action_space_size=5,
learning_rate=0.1, discount=0.99, epsilon=1.0):
self.state_bins = state_bins
self.action_space_size = action_space_size
self.lr = learning_rate
self.gamma = discount
self.epsilon = epsilon
self.epsilon_decay = 0.995
self.epsilon_min = 0.01
# Q-table
self.q_table = defaultdict(lambda: np.zeros(action_space_size))
# State discretization
self.state_bounds = [
(-10, 10), # x position
(-10, 10), # y position
(0, 20), # altitude
(-5, 5), # x velocity
(-5, 5), # y velocity
(-3, 3), # z velocity
]
def discretize_state(self, state):
discrete = []
for i, (s, (low, high)) in enumerate(zip(state, self.state_bounds)):
s_clipped = np.clip(s, low, high)
bins = np.linspace(low, high, self.state_bins)
discrete.append(np.digitize(s_clipped, bins) - 1)
return tuple(discrete)
def choose_action(self, state):
if np.random.random() < self.epsilon:
return np.random.randint(self.action_space_size)
discrete_state = self.discretize_state(state)
return np.argmax(self.q_table[discrete_state])
def learn(self, state, action, reward, next_state, done):
s = self.discretize_state(state)
s_next = self.discretize_state(next_state)
current_q = self.q_table[s][action]
if done:
target = reward
else:
target = reward + self.gamma * np.max(self.q_table[s_next])
self.q_table[s][action] += self.lr * (target - current_q)
def decay_epsilon(self):
self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)
class DiscreteActionSpace:
"""Discretized action space for Q-learning."""
ACTIONS = {
0: np.array([0, 0, 0]), # Hover
1: np.array([1, 0, 0]), # Forward
2: np.array([-1, 0, 0]), # Backward
3: np.array([0, 1, 0]), # Left
4: np.array([0, -1, 0]), # Right
5: np.array([0, 0, 1]), # Up
6: np.array([0, 0, -1]), # Down
7: np.array([1, 1, 0]), # Forward-Left
8: np.array([1, -1, 0]), # Forward-Right
9: np.array([-1, 1, 0]), # Backward-Left
10: np.array([-1, -1, 0]), # Backward-Right
}
def __init__(self):
self.n_actions = len(self.ACTIONS)
def get_action(self, idx):
return self.ACTIONS[idx]
Proximal Policy Optimization (PPO)
PPO is the most popular policy gradient method for continuous control. It's stable, sample-efficient, and works well for drone flight control.
PPO Architecture for Drones
PPO Implementation
python
from torch.distributions import Normal
class ActorCritic(nn.Module):
def __init__(self, state_dim, action_dim, hidden_size=256):
super().__init__()
# Shared encoder
self.encoder = nn.Sequential(
nn.Linear(state_dim, hidden_size),
nn.LayerNorm(hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.LayerNorm(hidden_size),
nn.ReLU()
)
# Actor (policy)
self.actor_mean = nn.Sequential(
nn.Linear(hidden_size, hidden_size // 2),
nn.ReLU(),
nn.Linear(hidden_size // 2, action_dim)
)
self.actor_log_std = nn.Parameter(torch.zeros(action_dim))
# Critic (value)
self.critic = nn.Sequential(
nn.Linear(hidden_size, hidden_size // 2),
nn.ReLU(),
nn.Linear(hidden_size // 2, 1)
)
def forward(self, state):
features = self.encoder(state)
# Actor
action_mean = self.actor_mean(features)
action_std = self.actor_log_std.exp().expand_as(action_mean)
# Critic
value = self.critic(features)
return action_mean, action_std, value
def get_action(self, state):
action_mean, action_std, value = self.forward(state)
dist = Normal(action_mean, action_std)
action = dist.sample()
log_prob = dist.log_prob(action).sum(dim=-1)
return action, log_prob, value
def evaluate(self, state, action):
action_mean, action_std, value = self.forward(state)
dist = Normal(action_mean, action_std)
log_prob = dist.log_prob(action).sum(dim=-1)
entropy = dist.entropy().sum(dim=-1)
return log_prob, entropy, value
class PPOTrainer:
def __init__(self, model, lr=3e-4, gamma=0.99, gae_lambda=0.95,
clip_epsilon=0.2, entropy_coef=0.01, value_coef=0.5):
self.model = model
self.optimizer = torch.optim.Adam(model.parameters(), lr=lr)
self.gamma = gamma
self.gae_lambda = gae_lambda
self.clip_epsilon = clip_epsilon
self.entropy_coef = entropy_coef
self.value_coef = value_coef
def compute_gae(self, rewards, values, dones, next_value):
advantages = []
gae = 0
for t in reversed(range(len(rewards))):
if t == len(rewards) - 1:
next_val = next_value
else:
next_val = values[t + 1]
delta = rewards[t] + self.gamma * next_val * (1 - dones[t]) - values[t]
gae = delta + self.gamma * self.gae_lambda * (1 - dones[t]) * gae
advantages.insert(0, gae)
return torch.tensor(advantages, dtype=torch.float32)
def update(self, rollout_buffer):
states = torch.tensor(rollout_buffer['states'], dtype=torch.float32)
actions = torch.tensor(rollout_buffer['actions'], dtype=torch.float32)
old_log_probs = torch.tensor(rollout_buffer['log_probs'], dtype=torch.float32)
rewards = rollout_buffer['rewards']
dones = rollout_buffer['dones']
values = rollout_buffer['values']
# Compute GAE
with torch.no_grad():
_, _, next_value = self.model(states[-1:])
returns = self.compute_gae(rewards, values, dones, next_value.item())
advantages = returns - torch.tensor(values, dtype=torch.float32)
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
# PPO update epochs
total_loss = 0
for _ in range(10):
log_probs, entropy, new_values = self.model.evaluate(states, actions)
# Policy loss with clipping
ratio = (log_probs - old_log_probs).exp()
surr1 = ratio * advantages
surr2 = torch.clamp(ratio, 1 - self.clip_epsilon, 1 + self.clip_epsilon) * advantages
policy_loss = -torch.min(surr1, surr2).mean()
# Value loss
value_loss = F.mse_loss(new_values.squeeze(), returns)
# Entropy bonus
entropy_loss = -entropy.mean()
# Total loss
loss = (policy_loss +
self.value_coef * value_loss +
self.entropy_coef * entropy_loss)
self.optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(self.model.parameters(), 0.5)
self.optimizer.step()
total_loss += loss.item()
return total_loss / 10
Soft Actor-Critic (SAC)
Sponsored
SAC is an off-policy algorithm that maximizes both reward and policy entropy, encouraging exploration. It's particularly effective for drone control with continuous action spaces.
python
from torch.distributions import Normal
class SACAgent:
def __init__(self, state_dim, action_dim, hidden_size=256):
# Twin Q-networks
self.q1 = self._build_critic(state_dim, action_dim, hidden_size)
self.q2 = self._build_critic(state_dim, action_dim, hidden_size)
self.q1_target = self._build_critic(state_dim, action_dim, hidden_size)
self.q2_target = self._build_critic(state_dim, action_dim, hidden_size)
# Policy network
self.policy = self._build_policy(state_dim, action_dim, hidden_size)
# Automatic entropy tuning
self.target_entropy = -action_dim
self.log_alpha = torch.zeros(1, requires_grad=True)
# Optimizers
self.q1_optimizer = torch.optim.Adam(self.q1.parameters(), lr=3e-4)
self.q2_optimizer = torch.optim.Adam(self.q2.parameters(), lr=3e-4)
self.policy_optimizer = torch.optim.Adam(self.policy.parameters(), lr=3e-4)
self.alpha_optimizer = torch.optim.Adam([self.log_alpha], lr=3e-4)
def _build_critic(self, state_dim, action_dim, hidden_size):
return nn.Sequential(
nn.Linear(state_dim + action_dim, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, 1)
)
def _build_policy(self, state_dim, action_dim, hidden_size):
return nn.Sequential(
nn.Linear(state_dim, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, action_dim * 2) # mean and log_std
)
def select_action(self, state, evaluate=False):
state_tensor = torch.FloatTensor(state).unsqueeze(0)
output = self.policy(state_tensor)
mean, log_std = output.chunk(2, dim=-1)
log_std = torch.clamp(log_std, -20, 2)
if evaluate:
action = torch.tanh(mean)
else:
std = log_std.exp()
normal = Normal(mean, std)
x_t = normal.rsample()
action = torch.tanh(x_t)
return action.detach().numpy()[0]
def update(self, batch, alpha=0.2):
states, actions, rewards, next_states, dones = batch
# Update Q-networks
with torch.no_grad():
next_actions, next_log_probs = self.policy.sample(next_states)
q1_next = self.q1_target(torch.cat([next_states, next_actions], dim=-1))
q2_next = self.q2_target(torch.cat([next_states, next_actions], dim=-1))
q_next = torch.min(q1_next, q2_next) - alpha * next_log_probs
target_q = rewards + (1 - dones) * 0.99 * q_next
q1_loss = F.mse_loss(self.q1(torch.cat([states, actions], dim=-1)), target_q)
q2_loss = F.mse_loss(self.q2(torch.cat([states, actions], dim=-1)), target_q)
self.q1_optimizer.zero_grad()
q1_loss.backward()
self.q1_optimizer.step()
self.q2_optimizer.zero_grad()
q2_loss.backward()
self.q2_optimizer.step()
# Update policy
new_actions, log_probs = self.policy.sample(states)
q1_new = self.q1(torch.cat([states, new_actions], dim=-1))
q2_new = self.q2(torch.cat([states, new_actions], dim=-1))
q_new = torch.min(q1_new, q2_new)
policy_loss = (alpha * log_probs - q_new).mean()
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
# Update entropy coefficient
alpha_loss = -(self.log_alpha * (log_probs + self.target_entropy).detach()).mean()
self.alpha_optimizer.zero_grad()
alpha_loss.backward()
self.alpha_optimizer.step()
return self.log_alpha.exp().item()
Reward Shaping for Drones
Reward design is critical for learning effective drone behaviors. Poorly designed rewards lead to unintended behaviors.
Reward Components
Reward Function Implementation
python
from typing import Dict
class DroneRewardFunction:
def __init__(self, config: Dict = None):
config = config or {}
self.goal_reward = config.get('goal_reward', 100.0)
self.collision_penalty = config.get('collision_penalty', -1000.0)
self.distance_scale = config.get('distance_scale', 1.0)
self.altitude_target = config.get('altitude_target', 5.0)
self.max_speed = config.get('max_speed', 5.0)
self.battery_threshold = config.get('battery_threshold', 20.0)
self.prev_distance = None
self.prev_altitude = None
self.prev_velocity = None
def compute_reward(self, state: Dict) -> float:
reward = 0.0
# Goal reaching
if state.get('reached_goal', False):
return self.goal_reward
# Collision
if state.get('collision', False):
return self.collision_penalty
# Progress toward goal (dense reward)
current_distance = state['distance_to_goal']
if self.prev_distance is not None:
progress = self.prev_distance - current_distance
reward += progress * self.distance_scale
# Altitude maintenance
altitude = state['altitude']
altitude_error = abs(altitude - self.altitude_target)
reward -= altitude_error * 0.1
# Speed penalty (encourage moderate speed)
speed = state['speed']
if speed > self.max_speed:
reward -= (speed - self.max_speed) * 0.5
# Smooth control (penalize jerky movements)
if self.prev_velocity is not None:
velocity_change = np.linalg.norm(state['velocity'] - self.prev_velocity)
reward -= velocity_change * 0.01
# Battery warning
battery = state.get('battery', 100)
if battery < self.battery_threshold:
reward -= (self.battery_threshold - battery) * 0.5
# Survival bonus
reward += 0.1
# Update previous values
self.prev_distance = current_distance
self.prev_altitude = altitude
self.prev_velocity = state['velocity'].copy() if 'velocity' in state else None
return reward
def reset(self):
self.prev_distance = None
self.prev_altitude = None
self.prev_velocity = None
Sim-to-Real Transfer
Training in simulation is cheap and safe, but policies must transfer to real drones. Domain randomization helps bridge the reality gap.
python
from typing import Dict
class DomainRandomizer:
def __init__(self):
self.randomization_params = {
'wind': {'mean': 0, 'std': 0.5},
'mass': {'min': 0.9, 'max': 1.1},
'drag': {'min': 0.8, 'max': 1.2},
'sensor_noise': {'std': 0.02},
'motor_delay': {'min': 0, 'max': 0.02},
'gravity': {'min': 9.75, 'max': 9.85},
}
def randomize_environment(self) -> Dict:
params = {}
for key, values in self.randomization_params.items():
if 'mean' in values and 'std' in values:
params[key] = np.random.normal(values['mean'], values['std'])
elif 'min' in values and 'max' in values:
params[key] = np.random.uniform(values['min'], values['max'])
elif 'std' in values:
params[key] = np.random.normal(0, values['std'])
return params
def apply_randomization(self, sim_env, params: Dict):
if 'wind' in params:
sim_env.set_wind(params['wind'])
if 'mass' in params:
sim_env.set_drone_mass(params['mass'])
if 'drag' in params:
sim_env.set_drag_coefficient(params['drag'])
if 'gravity' in params:
sim_env.set_gravity(params['gravity'])
Training Pipeline
python
from collections import deque
class RLTrainingPipeline:
def __init__(self, env, agent, reward_fn, randomizer=None):
self.env = env
self.agent = agent
self.reward_fn = reward_fn
self.randomizer = randomizer
self.episode_rewards = deque(maxlen=100)
self.episode_lengths = deque(maxlen=100)
self.best_reward = -float('inf')
def train(self, n_episodes=1000, max_steps=500):
for episode in range(n_episodes):
state = self.env.reset()
self.reward_fn.reset()
if self.randomizer:
params = self.randomizer.randomize_environment()
self.randomizer.apply_randomization(self.env, params)
episode_reward = 0
episode_length = 0
for step in range(max_steps):
# Select action
action = self.agent.select_action(state)
# Step environment
next_state, done, info = self.env.step(action)
# Compute reward
reward_info = {
'distance_to_goal': info['distance_to_goal'],
'altitude': info['altitude'],
'velocity': info['velocity'],
'speed': info['speed'],
'battery': info['battery'],
'collision': info.get('collision', False),
'reached_goal': info.get('reached_goal', False)
}
reward = self.reward_fn.compute_reward(reward_info)
# Store transition
self.agent.store_transition(state, action, reward, next_state, done)
# Learn
if len(self.agent.buffer) > self.agent.batch_size:
self.agent.update()
episode_reward += reward
episode_length += 1
state = next_state
if done:
break
self.episode_rewards.append(episode_reward)
self.episode_lengths.append(episode_length)
avg_reward = np.mean(self.episode_rewards)
if episode % 10 == 0:
print(f"Episode {episode}: Reward={episode_reward:.1f}, "
f"Avg={avg_reward:.1f}, Length={episode_length}")
if avg_reward > self.best_reward:
self.best_reward = avg_reward
self.agent.save_model('best_model.pth')
Hands-On Project: Complete RL Training
Advertisement
python
class DroneFlightEnv:
def __init__(self):
self.position = np.zeros(3)
self.velocity = np.zeros(3)
self.orientation = np.zeros(3)
self.battery = 100.0
self.goal = np.array([10.0, 10.0, 5.0])
self.max_steps = 500
self.step_count = 0
# Dynamics parameters
self.dt = 0.1
self.mass = 1.0
self.gravity = 9.81
self.max_force = 15.0
def reset(self):
self.position = np.zeros(3)
self.velocity = np.zeros(3)
self.orientation = np.zeros(3)
self.battery = 100.0
self.step_count = 0
return self._get_state()
def step(self, action):
self.step_count += 1
# Apply action (thrust commands for 4 rotors)
forces = action * self.max_force
total_thrust = np.sum(forces)
# Simple dynamics
acceleration = np.array([0, 0, total_thrust / self.mass - self.gravity])
self.velocity += acceleration * self.dt
self.position += self.velocity * self.dt
# Battery consumption
self.battery -= 0.05 * np.sum(np.abs(action))
# Check termination
done = False
info = {
'collision': self._check_collision(),
'reached_goal': self._check_goal(),
'distance_to_goal': np.linalg.norm(self.position - self.goal),
'altitude': self.position[2],
'velocity': self.velocity.copy(),
'speed': np.linalg.norm(self.velocity),
'battery': self.battery
}
if info['collision'] or info['reached_goal'] or self.battery <= 0:
done = True
if self.step_count >= self.max_steps:
done = True
return self._get_state(), done, info
def _get_state(self):
return np.concatenate([
self.position,
self.velocity,
self.orientation,
[self.battery],
self.goal
])
def _check_collision(self):
return self.position[2] < 0 or np.linalg.norm(self.position[:2]) > 100
def _check_goal(self):
return np.linalg.norm(self.position - self.goal) < 1.0
# Complete training script
def train_drone_rl():
env = DroneFlightEnv()
state_dim = 13 # pos(3) + vel(3) + orient(3) + battery(1) + goal(3)
action_dim = 4 # 4 rotor commands
# Create PPO agent
model = ActorCritic(state_dim, action_dim)
agent = PPOTrainer(model)
# Training loop
n_episodes = 1000
for episode in range(n_episodes):
state = env.reset()
episode_reward = 0
for step in range(500):
state_tensor = torch.FloatTensor(state).unsqueeze(0)
action, log_prob, value = model.get_action(state_tensor)
next_state, done, info = env.step(action.detach().numpy()[0])
# Simple reward
reward = -0.1 # Survival
reward += info['distance_to_goal'] * 0.5 # Progress
if info['reached_goal']:
reward += 100
if info['collision']:
reward -= 100
episode_reward += reward
if done:
break
state = next_state
if episode % 10 == 0:
print(f"Episode {episode}: Total Reward = {episode_reward:.1f}")
if __name__ == "__main__":
train_drone_rl()
Key Takeaways
- Q-learning works for discretized drone control but scales poorly
- PPO is the go-to algorithm for continuous drone control
- SAC excels at exploration with entropy regularization
- Reward shaping is critical — design rewards that encourage desired behavior
- Sim-to-real transfer requires domain randomization for robust policies
- Always train with safety constraints and monitor for reward hacking