企业官网建设流程全解析

强化学习目标函数 J(θ) 实战：从公式推导到 PyTorch 3 行代码实现

强化学习的目标函数 J(θ) 是算法设计的核心，但很多开发者会在理论推导和工程实现之间遇到断层。本文将用可运行的代码展示如何将数学公式转化为高效的 PyTorch 实现，让你真正掌握强化学习的计算本质。

1. 目标函数的数学本质

在强化学习中，目标函数 J(θ) = E[G(τ)] 表示在策略 πθ 下轨迹回报的期望。这个看似简单的公式包含三个关键计算层次：

1. 轨迹生成：根据当前策略与环境交互产生状态-动作-奖励序列
2. 回报计算：对单条轨迹的奖励进行折扣求和 G(τ) = Σγᵗrₜ
3. 期望估计：通过多次采样逼近数学期望

# 数学公式的直观表达 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(logits=self.fc(x)) def discount_rewards(rewards, gamma=0.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(), lr=0.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()

这个实现包含了强化学习目标函数的所有关键要素：

1. 策略网络的参数化表示
2. 轨迹回报的折扣计算
3. 策略梯度的自动求导
4. 参数更新的优化过程

6. 高级话题延伸

自动微分验证： 可以通过数值微分验证梯度计算的正确性：

def numerical_gradient(f, x, eps=1e-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的自动微分机制，我们可以将复杂的数学推导转化为简洁高效的代码实现。记住，关键在于构建正确的计算图，让框架自动处理梯度计算，而开发者专注于策略设计和训练逻辑。