强化学习目标函数 J(θ) 实战:从公式推导到 PyTorch 3 行代码实现 强化学习目标函数 J(θ) 实战从公式推导到 PyTorch 3 行代码实现强化学习的目标函数 J(θ) 是算法设计的核心但很多开发者会在理论推导和工程实现之间遇到断层。本文将用可运行的代码展示如何将数学公式转化为高效的 PyTorch 实现让你真正掌握强化学习的计算本质。1. 目标函数的数学本质在强化学习中目标函数 J(θ) E[G(τ)] 表示在策略 πθ 下轨迹回报的期望。这个看似简单的公式包含三个关键计算层次轨迹生成根据当前策略与环境交互产生状态-动作-奖励序列回报计算对单条轨迹的奖励进行折扣求和 G(τ) Σγᵗrₜ期望估计通过多次采样逼近数学期望# 数学公式的直观表达 def J(theta): trajectories [generate_trajectory(theta) for _ in range(N)] returns [compute_discounted_return(tau) for tau in trajectories] return np.mean(returns)在实际实现时我们需要处理两个工程难题如何高效计算梯度 ∇J(θ)如何避免重复采样带来的计算浪费2. 梯度计算的魔法技巧策略梯度定理告诉我们可以直接对期望形式求导∇J(θ) E[∇logπ(a|s) * G(τ)]这个推导过程看似复杂但PyTorch的自动微分机制可以帮我们自动完成。关键在于构建正确的计算图# 蒙特卡洛策略梯度实现 def policy_gradient(states, actions, returns): log_probs policy_network(states).log_prob(actions) return -(log_probs * returns).mean() # 负号因为PyTorch默认最小化关键技巧使用.log_prob()获取动作的对数概率回报G(τ)作为权重不参与求导需.detach()最终目标是最小化负回报3. 三行核心实现结合自动微分和策略梯度定理我们得到最精简的实现# 完整的目标函数实现 def compute_objective(trajectories): states, actions, rewards unpack(trajectories) returns discount(rewards) # 计算折扣回报 log_probs policy(states).log_prob(actions) return (log_probs * returns).mean() # 这就是J(θ)!对应的梯度计算只需三行def update_step(trajectories): J compute_objective(trajectories) J.backward() # 自动计算梯度 optimizer.step()4. 工程实践中的优化技巧实际应用中我们需要考虑以下改进回报标准化returns (returns - returns.mean()) / (returns.std() 1e-8)并行采样# 使用多进程同时收集多条轨迹 with mp.Pool() as pool: trajectories pool.map(collect_episode, [policy]*num_workers)价值函数基准advantages returns - value_network(states) # 降低方差典型训练循环结构for epoch in range(epochs): # 数据收集阶段 trajectories collect_trajectories(policy) # 计算阶段 returns compute_returns(trajectories) loss -compute_objective(trajectories, returns) # 更新阶段 optimizer.zero_grad() loss.backward() optimizer.step()5. 完整案例CartPole实现下面是一个完整的训练实例使用PyTorch实现策略梯度算法import torch import gym class PolicyNet(torch.nn.Module): def __init__(self, state_dim, action_dim): super().__init__() self.fc torch.nn.Sequential( torch.nn.Linear(state_dim, 64), torch.nn.ReLU(), torch.nn.Linear(64, action_dim) ) def forward(self, x): return torch.distributions.Categorical(logitsself.fc(x)) def discount_rewards(rewards, gamma0.99): R 0 returns [] for r in reversed(rewards): R r gamma * R returns.insert(0, R) return torch.tensor(returns) env gym.make(CartPole-v1) policy PolicyNet(env.observation_space.shape[0], env.action_space.n) optimizer torch.optim.Adam(policy.parameters(), lr0.01) for episode in range(1000): state env.reset() states, actions, rewards [], [], [] # 收集轨迹 while True: dist policy(torch.FloatTensor(state)) action dist.sample().item() next_state, reward, done, _ env.step(action) states.append(state) actions.append(action) rewards.append(reward) state next_state if done: break # 计算目标 returns discount_rewards(rewards) log_probs policy(torch.FloatTensor(states)).log_prob(torch.tensor(actions)) loss -(log_probs * returns).mean() # 参数更新 optimizer.zero_grad() loss.backward() optimizer.step()这个实现包含了强化学习目标函数的所有关键要素策略网络的参数化表示轨迹回报的折扣计算策略梯度的自动求导参数更新的优化过程6. 高级话题延伸自动微分验证 可以通过数值微分验证梯度计算的正确性def numerical_gradient(f, x, eps1e-4): grad torch.zeros_like(x) for i in range(len(x)): x_plus x.clone() x_plus[i] eps x_minus x.clone() x_minus[i] - eps grad[i] (f(x_plus) - f(x_minus)) / (2*eps) return grad混合精度训练 使用AMP加速训练过程scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): loss -compute_objective(trajectories) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()强化学习的目标函数实现既需要理解其数学本质又要掌握工程实现的技巧。通过PyTorch的自动微分机制我们可以将复杂的数学推导转化为简洁高效的代码实现。记住关键在于构建正确的计算图让框架自动处理梯度计算而开发者专注于策略设计和训练逻辑。