LLM

import torch
import torch.nn as nn
import torch.optim as optim

# =========================
# 1. 模型与环境
# =========================
policy = PolicyNetwork()      # 策略模型 pi_theta(a|s)
value_fn = ValueNetwork()     # critic V(s)
optimizer_policy = optim.Adam(policy.parameters(), lr=3e-6)
optimizer_value  = optim.Adam(value_fn.parameters(), lr=1e-3)

# 超参数
clip_epsilon = 0.2
gamma = 0.99
lambda_gae = 0.95

# =========================
# 2. Rollout 收集样本
# =========================
# 假设我们有 batch_size 个环境 / prompt
states, actions, rewards, dones, log_probs_old = [], [], [], [], []

for s in env_batch:   # 伪代码：每个环境/问题
    state = s.reset()
    done = False
    traj_reward = 0
    while not done:
        # 策略采样动作
        dist = policy(state)
        action = dist.sample()
        log_prob = dist.log_prob(action)

        next_state, reward, done, info = env.step(action)

        # 保存轨迹
        states.append(state)
        actions.append(action)
        rewards.append(reward)
        dones.append(done)
        log_probs_old.append(log_prob)

        state = next_state

# =========================
# 3. 计算 Advantage (GAE)
# =========================
advantages = []
returns = []
gae = 0
next_value = 0
for t in reversed(range(len(rewards))):
    delta = rewards[t] + gamma * next_value * (1 - dones[t]) - value_fn(states[t])
    gae = delta + gamma * lambda_gae * (1 - dones[t]) * gae
    advantages.insert(0, gae)
    returns.insert(0, gae + value_fn(states[t]))
    next_value = value_fn(states[t])

advantages = torch.tensor(advantages)
returns = torch.tensor(returns)
log_probs_old = torch.tensor(log_probs_old)

# =========================
# 4. PPO surrogate loss
# =========================
for _ in range(ppo_epochs):
    dist = policy(states)
    log_probs = dist.log_prob(actions)
    ratio = (log_probs - log_probs_old).exp()   # pi_theta / pi_old

    # clipped surrogate objective
    surr1 = ratio * advantages
    surr2 = torch.clamp(ratio, 1.0 - clip_epsilon, 1.0 + clip_epsilon) * advantages
    policy_loss = -torch.min(surr1, surr2).mean()

    # value function loss
    value_loss = nn.MSELoss()(value_fn(states), returns)

    # optimize
    optimizer_policy.zero_grad()
    policy_loss.backward()
    optimizer_policy.step()

    optimizer_value.zero_grad()
    value_loss.backward()
    optimizer_value.step()

# 假设策略和旧策略已经定义好
policy = ...       # 任何策略模型，返回 logits
pi_old = ...       # 旧策略，复制 policy 参数

optimizer = torch.optim.Adam(policy.parameters(), lr=lr)

# 假设 reward_model 已经定义好
def reward_model(states, actions):
    return torch.rand(len(actions))  # 仅示意

# 训练循环（单步示意）
for batch in range(100):
    states = torch.randn(batch_size, state_dim)

    all_advantages, all_log_probs_old, all_actions = [], [], []

    for state in states:
        dist_old = pi_old.get_dist(state)
        actions = dist_old.sample((G,))
        rewards = reward_model(state.repeat(G,1), actions)

        baseline = rewards.mean()
        advantages = rewards - baseline

        all_advantages.append(advantages)
        all_log_probs_old.append(dist_old.log_prob(actions))
        all_actions.append(actions)

    all_advantages = torch.cat(all_advantages)
    all_log_probs_old = torch.cat(all_log_probs_old)
    all_actions = torch.cat(all_actions)

    # 新策略
    dist_new = policy.get_dist(states.repeat(G,1))
    log_probs_new = dist_new.log_prob(all_actions)

    # ratio 和 clip
    ratio = torch.exp(log_probs_new - all_log_probs_old)
    surr1 = ratio * all_advantages
    surr2 = torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps) * all_advantages
    policy_loss = -torch.min(surr1, surr2).mean()

    # KL penalty
    kl = torch.distributions.kl_divergence(dist_old, dist_new).mean()
    loss = policy_loss + kl_beta * kl

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()