策略梯度训练的稳定性救星用TRPO的KL约束实现PyTorch下的可靠更新1. 策略梯度方法的稳定性困境在强化学习实践中策略梯度方法如A2C、PPO常面临一个令人头疼的问题——训练过程的不稳定性。你可能经历过这样的场景模型在某个训练阶段突然崩溃回报曲线出现断崖式下跌之前所有的调参努力都付诸东流。这种不稳定性主要源于两个核心因素策略更新的步长问题当策略网络是深度模型时沿策略梯度方向更新参数可能导致步长过大使策略突然变得极差数据分布的偏移策略更新后新策略生成的数据分布与旧策略差异过大导致基于旧数据计算的梯度估计失效传统策略梯度方法就像没有安全带的赛车——速度快但风险极高。而TRPOTrust Region Policy Optimization通过引入KL散度约束为策略更新装上了安全带确保每次更新都是安全的小步快跑。2. TRPO的核心思想与数学基础2.1 信任区域方法TRPO的核心思想是在策略优化过程中定义一个信任区域在这个区域内我们可以相信策略改进的近似是有效的。具体来说它通过限制新旧策略之间的KL散度来确保更新后的策略不会偏离当前策略太远。数学上TRPO的优化目标可以表示为maximize θ [πθ(a|s)/πθ(a|s) * Aθ(s,a)] subject to [DKL(πθ(·|s) || πθ(·|s))] ≤ δ其中θ是当前策略参数θ是待优化的新策略参数Aθ(s,a)是优势函数δ是KL散度的约束阈值2.2 重要性采样与目标函数TRPO使用重要性采样技术来利用旧策略的数据评估新策略的性能。其目标函数可以分解为策略梯度部分[πθ(a|s)/πθ(a|s) * Aθ(s,a)]KL约束部分[DKL(πθ(·|s) || πθ(·|s))] ≤ δ这种形式允许我们在不生成新数据的情况下评估策略改进同时通过KL约束保证评估的可靠性。3. TRPO的PyTorch实现关键步骤3.1 网络结构设计首先我们需要定义策略网络和价值网络class PolicyNet(torch.nn.Module): def __init__(self, state_dim, hidden_dim, action_dim): super().__init__() self.fc1 torch.nn.Linear(state_dim, hidden_dim) self.fc2 torch.nn.Linear(hidden_dim, action_dim) def forward(self, x): x F.relu(self.fc1(x)) return F.softmax(self.fc2(x), dim1) class ValueNet(torch.nn.Module): def __init__(self, state_dim, hidden_dim): super().__init__() self.fc1 torch.nn.Linear(state_dim, hidden_dim) self.fc2 torch.nn.Linear(hidden_dim, 1) def forward(self, x): x F.relu(self.fc1(x)) return self.fc2(x)3.2 共轭梯度法求解TRPO的核心难点在于处理KL约束。我们使用共轭梯度法来近似求解带约束的优化问题def conjugate_gradient(self, grad, states, old_action_dists): x torch.zeros_like(grad) r grad.clone() p grad.clone() rdotr torch.dot(r, r) for i in range(10): # 共轭梯度迭代次数 Hp self.hessian_matrix_vector_product(states, old_action_dists, p) alpha rdotr / torch.dot(p, Hp) x alpha * p r - alpha * Hp new_rdotr torch.dot(r, r) if new_rdotr 1e-10: break beta new_rdotr / rdotr p r beta * p rdotr new_rdotr return x3.3 线性搜索确保稳定性在得到更新方向后TRPO通过线性搜索找到合适的步长def line_search(self, states, actions, advantage, old_log_probs, old_action_dists, max_vec): old_para torch.nn.utils.convert_parameters.parameters_to_vector( self.actor.parameters()) old_obj self.compute_surrogate_obj(states, actions, advantage, old_log_probs, self.actor) for i in range(15): # 线性搜索迭代次数 coef self.alpha**i new_para old_para coef * max_vec new_actor copy.deepcopy(self.actor) torch.nn.utils.convert_parameters.vector_to_parameters( new_para, new_actor.parameters()) new_action_dists torch.distributions.Categorical(new_actor(states)) kl_div torch.mean( torch.distributions.kl.kl_divergence(old_action_dists, new_action_dists)) new_obj self.compute_surrogate_obj(states, actions, advantage, old_log_probs, new_actor) if new_obj old_obj and kl_div self.kl_constraint: return new_para return old_para4. TRPO在连续动作空间中的实现对于连续动作空间的问题我们需要调整策略网络的输出为高斯分布class PolicyNetContinuous(torch.nn.Module): def __init__(self, state_dim, hidden_dim, action_dim): super().__init__() self.fc1 torch.nn.Linear(state_dim, hidden_dim) self.fc_mu torch.nn.Linear(hidden_dim, action_dim) self.fc_std torch.nn.Linear(hidden_dim, action_dim) def forward(self, x): x F.relu(self.fc1(x)) mu 2.0 * torch.tanh(self.fc_mu(x)) # 限制均值范围 std F.softplus(self.fc_std(x)) # 保证标准差为正 return mu, std对应的动作采样和概率计算也需要相应调整def take_action(self, state): state torch.tensor(np.array([state]), dtypetorch.float).to(self.device) mu, std self.actor(state) action_dist torch.distributions.Normal(mu, std) return action_dist.sample().cpu().numpy().flatten()5. 实际训练中的技巧与调参5.1 超参数设置建议根据实践经验以下超参数组合通常能取得不错的效果超参数离散动作空间建议值连续动作空间建议值γ (折扣因子)0.990.90-0.95λ (GAE参数)0.950.85-0.95KL约束阈值0.00050.0005-0.001线性搜索α0.50.3-0.5Critic学习率1e-21e-3-1e-25.2 优势估计的改进TRPO通常与GAEGeneralized Advantage Estimation结合使用可以更稳定地估计优势函数def calculate_gae(rewards, values, gamma0.99, lambda_0.95): deltas [r gamma * v_next - v for r, v, v_next in zip(rewards, values, values[1:] [0])] advantages [] advantage 0 for delta in reversed(deltas): advantage gamma * lambda_ * advantage delta advantages.append(advantage) advantages.reverse() return advantages5.3 训练监控与调试在训练过程中建议监控以下指标平均回报KL散度的实际值优势函数的均值和方差策略更新的接受率线性搜索成功次数当出现以下情况时可能需要调整超参数KL散度远低于约束阈值 → 可适当增大步长线性搜索频繁失败 → 减小KL约束阈值或调整初始步长回报波动剧烈 → 检查优势估计或减小学习率