Imitation Learning for Drone Control
Instead of designing reward functions, imitation learning lets drones learn by watching expert demonstrations. Think of it like teaching a child to ride a bike â they learn by observing and imitating, not by mathematical optimization.
Why Imitation Learning?
Sponsored
Traditional RL requires careful reward engineering. Imitation learning bypasses this by learning directly from expert behavior. For drones, experts can be human pilots, classical controllers, or other RL-trained agents.
Imitation Learning Pipeline
[SVG Diagram - see original for visual]
Behavior Cloning (BC)
Behavior cloning is the simplest form of imitation learning â just supervised learning on state-action pairs from expert demonstrations.
Behavior Cloning Implementation
python
from typing import List, Tuple
class BehaviorCloningNetwork(nn.Module):
def __init__(self, state_dim: int, action_dim: int, hidden_size: int = 256):
super().__init__()
self.network = nn.Sequential(
nn.Linear(state_dim, hidden_size),
nn.ReLU(),
nn.BatchNorm1d(hidden_size),
nn.Dropout(0.1),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.BatchNorm1d(hidden_size),
nn.Dropout(0.1),
nn.Linear(hidden_size, hidden_size // 2),
nn.ReLU(),
nn.Linear(hidden_size // 2, action_dim)
)
def forward(self, state):
return self.network(state)
class BehaviorCloningTrainer:
def __init__(self, model: nn.Module, lr: float = 1e-3):
self.model = model
self.optimizer = torch.optim.Adam(model.parameters(), lr=lr)
self.loss_fn = nn.MSELoss()
def train_epoch(self, dataloader) -> float:
self.model.train()
total_loss = 0
for states, actions in dataloader:
predicted_actions = self.model(states)
loss = self.loss_fn(predicted_actions, actions)
self.optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
self.optimizer.step()
total_loss += loss.item()
return total_loss / len(dataloader)
def evaluate(self, dataloader) -> float:
self.model.eval()
total_loss = 0
with torch.no_grad():
for states, actions in dataloader:
predicted_actions = self.model(states)
loss = self.loss_fn(predicted_actions, actions)
total_loss += loss.item()
return total_loss / len(dataloader)
class DemonstrationDataset:
def __init__(self):
self.states = []
self.actions = []
def add_trajectory(self, states: np.ndarray, actions: np.ndarray):
self.states.extend(states)
self.actions.extend(actions)
def to_tensors(self) -> Tuple[torch.Tensor, torch.Tensor]:
states = torch.FloatTensor(np.array(self.states))
actions = torch.FloatTensor(np.array(self.actions))
return states, actions
def __len__(self):
return len(self.states)
def __getitem__(self, idx):
return self.states[idx], self.actions[idx]
Distribution Shift Problem
Behavior cloning suffers from distribution shift â the trained policy encounters states not seen during training, leading to compounding errors.
python
from typing import Optional
class DistributionShiftAnalyzer:
def __init__(self, expert_policy, learned_policy):
self.expert = expert_policy
self.learned = learned_policy
def analyze_drift(self, initial_states: np.ndarray, n_steps: int = 50):
expert_trajectories = []
learned_trajectories = []
for state in initial_states:
expert_traj = self._rollout(self.expert, state, n_steps)
learned_traj = self._rollout(self.learned, state, n_steps)
expert_trajectories.append(expert_traj)
learned_trajectories.append(learned_traj)
# Compute state distribution divergence
drift = self._compute_drift(learned_trajectories, expert_trajectories)
return {
'expert_states': expert_trajectories,
'learned_states': learned_trajectories,
'state_drift': drift
}
def _rollout(self, policy, initial_state, n_steps):
states = [initial_state]
state = initial_state
for _ in range(n_steps):
action = policy.predict(state)
state = self._sim_step(state, action)
states.append(state)
return np.array(states)
def _sim_step(self, state, action):
# Simple simulation
next_state = state.copy()
next_state[:3] += action[:3] * 0.1
return next_state
def _compute_drift(self, learned_trajs, expert_trajs):
drifts = []
for learned, expert in zip(learned_trajs, expert_trajs):
for t in range(min(len(learned), len(expert))):
distance = np.linalg.norm(learned[t] - expert[t])
drifts.append(distance * t) # Weight by time step
return np.mean(drifts)
DAgger (Dataset Aggregation)
DAgger addresses distribution shift by iteratively collecting new data from the learned policy and labeling it with expert actions.
DAgger Algorithm
[SVG Diagram - see original for visual]
DAgger Implementation
python
from typing import Callable
class DAgger:
def __init__(self, expert_policy: Callable, student_model: nn.Module,
env, lr: float = 1e-3, beta: float = 0.5):
self.expert = expert_policy
self.student = student_model
self.env = env
self.beta = beta # Interpolation factor for expert data
self.optimizer = torch.optim.Adam(student_model.parameters(), lr=lr)
self.dataset = {'states': [], 'expert_actions': []}
# Initial training with expert demonstrations
self._initial_dataset = None
def initial_train(self, expert_states, expert_actions, n_epochs=50):
"""Train on initial expert demonstrations."""
self._initial_dataset = (expert_states, expert_actions)
states = torch.FloatTensor(expert_states)
actions = torch.FloatTensor(expert_actions)
for epoch in range(n_epochs):
pred = self.student(states)
loss = nn.MSELoss()(pred, actions)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
def dagger_iteration(self, n_episodes=5, n_steps=200):
"""Run one DAgger iteration."""
new_states = []
new_expert_actions = []
for _ in range(n_episodes):
state = self.env.reset()
for step in range(n_steps):
# Student chooses action
state_tensor = torch.FloatTensor(state).unsqueeze(0)
with torch.no_grad():
student_action = self.student(state_tensor).numpy()[0]
# With probability (1-beta), use expert action
if np.random.random() < (1 - self.beta):
expert_action = self.expert(state)
else:
expert_action = student_action
new_states.append(state)
new_expert_actions.append(expert_action)
# Step environment
state, done, _ = self.env.step(student_action)
if done:
break
# Add new data to dataset
self.dataset['states'].extend(new_states)
self.dataset['expert_actions'].extend(new_expert_actions)
# Retrain on aggregated dataset
self._retrain()
return len(new_states)
def _retrain(self, n_epochs=30):
states = torch.FloatTensor(np.array(self.dataset['states']))
actions = torch.FloatTensor(np.array(self.dataset['expert_actions']))
for epoch in range(n_epochs):
indices = torch.randperm(len(states))
for i in range(0, len(states), 32):
batch_idx = indices[i:i+32]
batch_states = states[batch_idx]
batch_actions = actions[batch_idx]
pred = self.student(batch_states)
loss = nn.MSELoss()(pred, batch_actions)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
def train(self, n_iterations=10, n_episodes_per_iter=5):
"""Run full DAgger training."""
for i in range(n_iterations):
n_samples = self.dagger_iteration(n_episodes_per_iter)
print(f"Iteration {i+1}: Added {n_samples} samples, "
f"Total dataset: {len(self.dataset['states'])}")
Learning from Human Pilots
Advertisement
Human pilots provide natural, intuitive demonstrations that capture complex flying behaviors.
python
from datetime import datetime
from typing import List, Dict
class HumanDemonstrationCollector:
def __init__(self, control_interface):
self.interface = control_interface
self.recordings = []
self.current_recording = None
self.is_recording = False
def start_recording(self, metadata: Dict = None):
self.current_recording = {
'timestamp': datetime.now().isoformat(),
'metadata': metadata or {},
'states': [],
'actions': [],
'timestamps': []
}
self.is_recording = True
print("Recording started...")
def record_step(self, state: np.ndarray, action: np.ndarray):
if self.is_recording and self.current_recording:
self.current_recording['states'].append(state.tolist())
self.current_recording['actions'].append(action.tolist())
self.current_recording['timestamps'].append(
datetime.now().timestamp()
)
def stop_recording(self):
self.is_recording = False
if self.current_recording:
self.recordings.append(self.current_recording)
print(f"Recording saved: {len(self.current_recording['states'])} steps")
self.current_recording = None
def save_demonstrations(self, filepath: str):
data = {
'demonstrations': self.recordings,
'total_demos': len(self.recordings),
'total_steps': sum(len(r['states']) for r in self.recordings)
}
with open(filepath, 'w') as f:
json.dump(data, f, indent=2)
print(f"Saved {len(self.recordings)} demonstrations to {filepath}")
def load_demonstrations(self, filepath: str):
with open(filepath, 'r') as f:
data = json.load(f)
self.recordings = data['demonstrations']
print(f"Loaded {len(self.recordings)} demonstrations")
def filter_by_quality(self, min_steps=50, max_jerk=10.0):
filtered = []
for recording in self.recordings:
if len(recording['states']) < min_steps:
continue
# Compute smoothness
actions = np.array(recording['actions'])
if len(actions) > 1:
jerk = np.mean(np.abs(np.diff(actions, axis=0)))
if jerk > max_jerk:
continue
filtered.append(recording)
self.recordings = filtered
print(f"Kept {len(filtered)} demonstrations after quality filter")
class DataAugmentation:
"""Augment drone demonstration data."""
@staticmethod
def add_noise(states, actions, noise_std=0.05):
noisy_states = states + np.random.normal(0, noise_std, states.shape)
return noisy_states, actions
@staticmethod
def time_warp(states, actions, warping_factor=0.2):
n = len(states)
warp = 1 + np.random.uniform(-warping_factor, warping_factor)
indices = np.clip((np.arange(n) * warp).astype(int), 0, n - 1)
return states[indices], actions[indices]
@staticmethod
def rotate_demonstration(states, actions, max_angle=15):
angle = np.random.uniform(-max_angle, max_angle) * np.pi / 180
cos_a, sin_a = np.cos(angle), np.sin(angle)
rotation = np.array([[cos_a, -sin_a], [sin_a, cos_a]])
augmented_states = states.copy()
augmented_states[:, :2] = states[:, :2] @ rotation.T
augmented_states[:, 5] += angle # Yaw
return augmented_states, actions
Generative Adversarial Imitation Learning (GAIL)
GAIL uses adversarial training to match the state-action distribution of the expert without requiring paired state-action labels.
python
class GAILDiscriminator(nn.Module):
def __init__(self, state_dim, action_dim, hidden_size=128):
super().__init__()
self.network = nn.Sequential(
nn.Linear(state_dim + action_dim, hidden_size),
nn.LeakyReLU(0.2),
nn.Linear(hidden_size, hidden_size),
nn.LeakyReLU(0.2),
nn.Linear(hidden_size, 1),
nn.Sigmoid()
)
def forward(self, state_action):
return self.network(state_action)
class GAIL:
def __init__(self, policy, expert_states, expert_actions,
state_dim, action_dim, lr=1e-3):
self.policy = policy
self.discriminator = GAILDiscriminator(state_dim, action_dim)
self.disc_optimizer = torch.optim.Adam(self.discriminator.parameters(), lr=lr)
# Expert data
self.expert_states = torch.FloatTensor(expert_states)
self.expert_actions = torch.FloatTensor(expert_actions)
self.expert_sa = torch.cat([self.expert_states, self.expert_actions], dim=-1)
def train_discriminator(self, policy_states, policy_actions, n_epochs=10):
policy_sa = torch.cat([policy_states, policy_actions], dim=-1)
for _ in range(n_epochs):
# Expert: label 1, Policy: label 0
expert_pred = self.discriminator(self.expert_sa)
policy_pred = self.discriminator(policy_sa.detach())
expert_loss = -torch.mean(torch.log(expert_pred + 1e-8))
policy_loss = -torch.mean(torch.log(1 - policy_pred + 1e-8))
loss = (expert_loss + policy_loss) / 2
self.disc_optimizer.zero_grad()
loss.backward()
self.disc_optimizer.step()
def compute_reward(self, state, action):
sa = torch.FloatTensor(np.concatenate([state, action]))
with torch.no_grad():
D = self.discriminator(sa.unsqueeze(0))
reward = -torch.log(1 - D + 1e-8)
return reward.item()
Hands-On Project: Complete Imitation Learning System
python
from typing import List, Tuple
class CompleteImitationLearningSystem:
def __init__(self, state_dim, action_dim):
self.state_dim = state_dim
self.action_dim = action_dim
# Models
self.bc_model = BehaviorCloningNetwork(state_dim, action_dim)
self.dagger_model = BehaviorCloningNetwork(state_dim, action_dim)
# Trainers
self.bc_trainer = BehaviorCloningTrainer(self.bc_model)
self.dagger = DAgger(
expert_policy=self._expert_policy,
student_model=self.dagger_model,
env=None # Set during training
)
# Data
self.demonstrations = DemonstrationDataset()
def _expert_policy(self, state):
"""Simulated expert policy (PD controller)."""
target = np.array([10.0, 10.0, 5.0])
error = target - state[:3]
kp = 1.0
kd = 0.5
action = kp * error - kd * state[3:6]
return np.clip(action, -1, 1)
def collect_demonstrations(self, n_trajectories=100):
print(f"Collecting {n_trajectories} expert demonstrations...")
for i in range(n_trajectories):
states = []
actions = []
state = np.random.uniform(-5, 5, self.state_dim)
for step in range(200):
action = self._expert_policy(state)
states.append(state)
actions.append(action)
# Simulate step
state = self._sim_step(state, action)
self.demonstrations.add_trajectory(
np.array(states),
np.array(actions)
)
print(f"Collected {len(self.demonstrations)} state-action pairs")
def _sim_step(self, state, action):
next_state = state.copy()
next_state[:3] += state[3:6] * 0.1 + 0.5 * action * 0.01
next_state[3:6] += action * 0.1
return next_state
def train_behavior_cloning(self, epochs=50):
print("\nTraining Behavior Cloning...")
states, actions = self.demonstrations.to_tensors()
dataset = torch.utils.data.TensorDataset(states, actions)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)
for epoch in range(epochs):
loss = self.bc_trainer.train_epoch(dataloader)
if epoch % 10 == 0:
print(f" Epoch {epoch}: Loss = {loss:.4f}")
def train_dagger(self, n_iterations=5, env=None):
print("\nTraining with DAgger...")
self.dagger.env = env
for i in range(n_iterations):
n_samples = self.dagger.dagger_iteration(n_episodes=5)
print(f" Iteration {i+1}: Added {n_samples} samples")
def evaluate(self, test_states, use_dagger=True):
model = self.dagger_model if use_dagger else self.bc_model
model.eval()
states_tensor = torch.FloatTensor(test_states)
with torch.no_grad():
predicted_actions = model(states_tensor).numpy()
expert_actions = np.array([self._expert_policy(s) for s in test_states])
mse = np.mean((predicted_actions - expert_actions) ** 2)
return mse
def _sim_step(self, state, action):
next_state = state.copy()
next_state[:3] += state[3:6] * 0.1 + 0.5 * action * 0.01
next_state[3:6] += action * 0.1
return next_state
# Usage example
def main():
system = CompleteImitationLearningSystem(state_dim=6, action_dim=3)
# Collect expert demonstrations
system.collect_demonstrations(n_trajectories=50)
# Train behavior cloning
system.train_behavior_cloning(epochs=50)
# Evaluate
test_states = np.random.uniform(-5, 5, (100, 6))
mse_bc = system.evaluate(test_states, use_dagger=False)
print(f"\nBehavior Cloning MSE: {mse_bc:.4f}")
if __name__ == "__main__":
main()
Key Takeaways
- Behavior cloning is simple but suffers from distribution shift
- DAgger iteratively improves by labeling learner's own trajectories
- GAIL enables imitation without explicit action labels
- Data quality matters more than quantity for imitation learning
- Combining multiple methods often yields best results
- Human demonstrations capture intuitive flying skills hard to code