【Agentic RL / 强化学习 / OPD】OpenClaw-RL 源码阅读笔记 — (2)— On-Policy Distillation

一、引言：从策略蒸馏到 On-Policy Distillation在强化学习中，策略蒸馏（Policy Distillation）是一种将复杂教师模型的知识迁移到轻量学生模型的技术。传统蒸馏通常使用离线数据，但存在分布偏移问题——学生模型在训练时遇到的状态与教师模型生成的状态不一致，导致性能下降。On-Policy Distillation（OPD） 通过让学生在教师策略的指导下进行在线交互，实时对齐状态分布，从而解决该问题。本文将结合 OpenClaw-RL 源码，从基础概念到高级实现，逐步解析 OPD 的核心原理。你将学习：- 什么是策略蒸馏及 On-Policy 变体- 如何用代码实现一个简单的 OPD 训练循环- OpenClaw-RL 中 OPD 的工程化设计## 二、基础概念：策略蒸馏的核心思想### 2.1 教师-学生框架在强化学习中，教师策略 ${\pi }_{teacher}(a|s)$ 通常是一个训练充分、性能优异的大模型（如深层神经网络），而学生策略 ${\pi }_{student}(a|s)$ 是一个轻量模型。蒸馏的目标是让学生模仿教师的动作分布，从而继承教师的性能。数学表达： 最小化 KL 散度： $${\mathcal{L}}_{distill}={\mathbb{E}}_{s\sim \mathcal{D}}\left[D_{KL}({\pi }_{teacher}(\cdot |s)\parallel {\pi }_{student}(\cdot |s))\right]$$其中 $\mathcal{D}$ 是训练数据的状态分布。### 2.2 On-Policy 蒸馏的动机传统蒸馏使用离线数据（教师与环境交互产生的固定数据集），但学生模型在部署时会遇到新状态，导致分布偏移。On-Policy Distillation 动态生成数据：学生与环境交互，同时教师提供动作指导，使状态分布与学生当前策略一致。## 三、代码示例1：实现基础 On-Policy 蒸馏循环我们用一个简单环境（如 CartPole）演示 OPD 的核心逻辑。假设教师策略已训练好。pythonimport torchimport torch.nn as nnimport torch.optim as optimimport gymfrom torch.distributions import Categorical# 定义简单的策略网络class PolicyNetwork(nn.Module): def __init__(self, state_dim, action_dim): super().__init__() self.fc1 = nn.Linear(state_dim, 64) self.fc2 = nn.Linear(64, action_dim) def forward(self, x): x = torch.relu(self.fc1(x)) return torch.softmax(self.fc2(x), dim=-1)# 假设已经加载了教师模型（这里用随机初始化模拟）teacher = PolicyNetwork(4, 2) # CartPole: state_dim=4, action_dim=2student = PolicyNetwork(4, 2)# 优化器optimizer = optim.Adam(student.parameters(), lr=0.001)# 环境env = gym.make("CartPole-v1")# On-Policy 蒸馏训练循环num_episodes = 100for episode in range(num_episodes): state = env.reset() done = False total_loss = 0.0 while not done: # 将状态转为张量 state_tensor = torch.FloatTensor(state).unsqueeze(0) # 学生策略采样动作（与环境交互） student_probs = student(state_tensor) student_dist = Categorical(student_probs) action = student_dist.sample().item() # 教师提供指导：计算教师动作概率（作为目标分布） with torch.no_grad(): teacher_probs = teacher(state_tensor) # 教师输出 softmax 概率 # 计算蒸馏损失：KL散度 教师 || 学生 # 注意：这里使用教师的概率作为目标，学生作为预测 loss = torch.sum(teacher_probs * (torch.log(teacher_probs + 1e-8) - torch.log(student_probs + 1e-8))) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step() # 环境步进 next_state, reward, done, _ = env.step(action) state = next_state total_loss += loss.item() print(f"Episode {episode}, Loss: {total_loss:.4f}")env.close()关键点：- 学生用当前策略与环境交互（On-Policy）- 教师不与环境交互，仅提供动作概率指导- 损失函数直接作用于状态-动作对的概率分布## 四、高级实现：OpenClaw-RL 中的 OPD 设计OpenClaw-RL 将 OPD 模块化，融入多智能体训练框架。其核心设计包括：### 4.1 数据流与并行化OpenClaw-RL 使用 Vectorized Environments 并行收集数据。每个环境中的智能体使用学生策略，但教师策略作为“影子模型”实时提供目标分布。### 4.2 关键组件：DistillationBuffer为减少样本相关性，OpenClaw-RL 设计了专用缓冲区，存储（状态，教师概率，学生动作）三元组。### 4.3 代码示例2：OpenClaw-RL 风格的 OPD 训练以下代码模拟 OpenClaw-RL 的 OPD 训练逻辑，包含缓冲区与分布式更新。pythonimport numpy as npimport torchimport torch.nn as nnimport torch.optim as optimfrom collections import deque# 假设环境为连续控制（简化版）class SimpleEnv: def __init__(self): self.state_dim = 3 self.action_dim = 2 def reset(self): return np.random.randn(self.state_dim) def step(self, action): # 模拟环境动态 next_state = np.random.randn(self.state_dim) reward = np.random.random() done = np.random.random() < 0.1 return next_state, reward, done, {}# 定义蒸馏缓冲区class DistillationBuffer: def __init__(self, capacity=1000): self.buffer = deque(maxlen=capacity) def add(self, state, teacher_probs, student_action): self.buffer.append((state, teacher_probs, student_action)) def sample(self, batch_size): indices = np.random.choice(len(self.buffer), batch_size, replace=False) batch = [self.buffer[i] for i in indices] states = torch.FloatTensor([b[0] for b in batch]) teacher_probs = torch.FloatTensor([b[1] for b in batch]) student_actions = torch.LongTensor([b[2] for b in batch]) return states, teacher_probs, student_actions# 初始化env = SimpleEnv()teacher = nn.Linear(3, 2) # 简化教师模型student = nn.Linear(3, 2)optimizer = optim.Adam(student.parameters(), lr=0.001)buffer = DistillationBuffer(capacity=500)# 训练循环（模拟 On-Policy 特性）num_iters = 200batch_size = 32for iteration in range(num_iters): # 收集 on-policy 数据 state = env.reset() done = False while not done: state_tensor = torch.FloatTensor(state).unsqueeze(0) # 学生采样动作 student_logits = student(state_tensor) student_probs = torch.softmax(student_logits, dim=-1) action = torch.multinomial(student_probs, 1).item() # 教师提供概率（无梯度） with torch.no_grad(): teacher_logits = teacher(state_tensor) teacher_probs = torch.softmax(teacher_logits, dim=-1) # 存储到缓冲区 buffer.add(state, teacher_probs.squeeze(0).numpy(), action) # 环境步进 next_state, reward, done, _ = env.step(action) state = next_state # 从缓冲区采样并更新学生（离线更新，但数据是 on-policy 收集的） if len(buffer.buffer) >= batch_size: states, teacher_probs, student_actions = buffer.sample(batch_size) # 计算蒸馏损失：交叉熵 + KL 散度混合 student_logits = student(states) student_probs = torch.softmax(student_logits, dim=-1) # 动作一致性损失（可选） ce_loss = nn.CrossEntropyLoss()(student_logits, student_actions) # KL 散度损失 kl_loss = torch.mean(torch.sum(teacher_probs * (torch.log(teacher_probs + 1e-8) - torch.log(student_probs + 1e-8)), dim=-1)) # 总损失（可调节权重） total_loss = 0.5 * ce_loss + 0.5 * kl_loss optimizer.zero_grad() total_loss.backward() optimizer.step() if iteration % 50 == 0: print(f"Iter {iteration}, Loss: {total_loss.item():.4f}")工程化特点：- 数据复用：缓冲区存储 on-policy 数据，支持小批量更新- 混合损失：结合动作一致性（交叉熵）与分布匹配（KL 散度）- 可扩展性：支持多智能体并行数据收集## 五、深入分析：OPD 的优势与挑战### 5.1 优势1. 分布对齐：学生策略生成的状态分布与训练数据一致，避免离线蒸馏的分布偏移2. 样本效率：在线交互反馈能快速纠正学生错误3. 渐进式学习：学生能力提升后，能探索更复杂状态，教师持续指导### 5.2 挑战1. 计算成本：每个时间步需计算教师模型的前向传播（可缓存或异步更新）2. 教师依赖性：教师策略质量直接影响学生上限3. 收敛稳定性：需谨慎调整蒸馏权重与环境奖励的平衡## 六、总结本文从策略蒸馏的基础概念出发，逐步深入到 On-Policy Distillation 的实现。通过两个可运行代码示例，我们展示了：1. 基础 OPD 训练循环：学生在线交互，教师提供概率指导2. OpenClaw-RL 工程化设计：缓冲区、混合损失、并行数据收集OPD 通过在线交互解决了分布偏移问题，在机器人控制（如 OpenClaw）和多智能体系统中具有重要应用。理解其源码设计，能帮助我们更好地部署强化学习模型到实际场景中。下一期将探讨 OPD 与 PPO 算法的结合，敬请期待。