AI强化学习实战指南：从策略梯度到DPO对齐

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

env = gym.make('CartPole-v1')
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n

class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=128):
        super().__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, action_dim)

    def forward(self, state):
        x = torch.relu(self.fc1(state))
        return torch.softmax(self.fc2(x), dim=-1)

def reinforce(num_episodes=1000, gamma=0.99, lr=1e-3):
    policy = PolicyNetwork(state_dim, action_dim)
    optimizer = optim.Adam(policy.parameters(), lr=lr)

    for episode in range(num_episodes):
        state = env.reset()
        log_probs = []
        rewards = []

        # 采样一条轨迹
        done = False
        while not done:
            state_tensor = torch.FloatTensor(state)
            action_probs = policy(state_tensor)
            action_dist = torch.distributions.Categorical(action_probs)
            action = action_dist.sample()

            next_state, reward, done, _ = env.step(action.item())

            log_probs.append(action_dist.log_prob(action))
            rewards.append(reward)
            state = next_state

        # 计算折扣回报
        returns = []
        R = 0
        for r in reversed(rewards):
            R = r + gamma * R
            returns.insert(0, R)

        # 标准化
        returns = torch.FloatTensor(returns)
        returns = (returns - returns.mean()) / (returns.std() + 1e-9)

        # 计算损失
        policy_loss = []
        for log_prob, R in zip(log_probs, returns):
            policy_loss.append(-log_prob * R)

        optimizer.zero_grad()
        policy_loss = torch.stack(policy_loss).sum()
        policy_loss.backward()
        optimizer.step()

# 糟糕的设计
reward = 1 if score > previous_score else -1

# 改进的设计
reward = base_reward + 0.1 * (current_score - previous_score)

action_probs = policy(state_tensor)
action_probs = action_probs * (1 + epsilon * torch.randn_like(action_probs))
action_probs = torch.softmax(action_probs, dim=-1)
epsilon = max(0.01, epsilon * 0.995)  # 衰减

# 监控梯度
for name, param in policy.named_parameters():
    if param.grad is not None:
        print(f"{name} grad norm: {param.grad.norm().item()}")

torch.nn.utils.clip_grad_norm_(policy.parameters(), max_norm=1.0)

class ActorCritic(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=128):
        super().__init__()
        # Actor：决策
        self.actor = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim)
        )
        # Critic：估值
        self.critic = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )

    def forward(self, state):
        action_probs = torch.softmax(self.actor(state), dim=-1)
        value = self.critic(state)
        return action_probs, value

# 数据格式
preference_data = [
    {
        "prompt": "用户：什么是量子计算？",
        "chosen": "量子计算利用量子力学原理...",
        "rejected": "量子计算很复杂，涉及薛定谔方程..."
    },
    # ...
]

def clean_preference_data(raw_data):
    cleaned = []
    for item in raw_data:
        # 过滤长度差异过大的对比
        len_diff = abs(len(item['chosen']) - len(item['rejected']))
        if len_diff > len(item['chosen']) * 0.5:
            continue
        # 检查标注一致性
        if item['chosen'] == item['rejected']:
            continue
        cleaned.append(item)
    return cleaned

class RewardModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.base = base_model
        self.score_head = nn.Linear(base_model.config.hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        outputs = self.base(input_ids=input_ids, attention_mask=attention_mask)
        last_hidden = outputs.last_hidden_state[:, 0, :]
        return self.score_head(last_hidden)

def compute_loss_with_regularization(chosen_scores, rejected_scores,
                                      chosen_input_ids, rejected_input_ids,
                                      lambda_reg=0.01):
    # 标准偏好损失
    loss = -torch.mean(torch.log(torch.sigmoid(chosen_scores - rejected_scores)))

    # 长度正则化（防止偏好特定长度）
    chosen_lengths = (chosen_input_ids != tokenizer.pad_token_id).sum(dim=1)
    rejected_lengths = (rejected_input_ids != tokenizer.pad_token_id).sum(dim=1)
    length_penalty = lambda_reg * torch.mean(torch.abs(chosen_lengths - rejected_lengths))

    return loss + length_penalty

def ppo_loss(log_probs, old_log_probs, advantages, clip_epsilon=0.2):
    ratio = torch.exp(log_probs - old_log_probs)
    clipped_ratio = torch.clamp(ratio, 1 - clip_epsilon, 1 + clip_epsilon)

    unclipped_loss = -ratio * advantages
    clipped_loss = -clipped_ratio * advantages
    return torch.mean(torch.maximum(unclipped_loss, clipped_loss))

def adjusted_reward(original_reward, generated_text, reference_text):
    base_reward = original_reward

    # 长度约束
    answer_length = len(generated_text.split())
    if answer_length < 20 or answer_length > 500:
        base_reward -= 0.5

    # 重复性惩罚
    ngrams = set()
    for i in range(len(generated_text.split()) - 2):
        ngram = ' '.join(generated_text.split()[i:i+3])
        if ngram in ngrams:
            base_reward -= 0.1
        ngrams.add(ngram)

    return base_reward

from trl import DPOTrainer, DPOConfig
from datasets import Dataset

train_dataset = Dataset.from_list(preference_data)

dpo_config = DPOConfig(
    output_dir="./dpo_model",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    learning_rate=5e-7,      # 比 RLHF 小很多
    num_train_epochs=3,      # 2-3 轮够
    beta=0.1,                # KL 散度系数
    max_length=512,
    max_prompt_length=256,
    bf16=True,
    gradient_checkpointing=True,
)

dpo_trainer = DPOTrainer(
    model=model,             # 必须是 SFT 过的模型
    ref_model=None,          # 自动创建参考副本
    args=dpo_config,
    train_dataset=train_dataset,
    tokenizer=tokenizer,
    beta=0.1,
)

dpo_trainer.train()

# 训练日志
{
    "train_loss": 0.456,
    "rewards/chosen": 0.789,    # 选中回答的隐含奖励
    "rewards/rejected": -0.456, # 拒绝回答的隐含奖励
    "rewards/acc": 0.85,        # 偏好预测准确率
    "rewards/margins": 1.245,   # 奖励边际
}

dpo_config = DPOConfig(
    gradient_checkpointing=True,
    optim="adamw_torch_fused",
    bf16=True,
    # 或者用 DeepSpeed
    # deepspeed="ds_config.json",
)

dpo_config = DPOConfig(
    learning_rate=1e-7,         # 降低
    warmup_ratio=0.1,           # warmup
    lr_scheduler_type="cosine", # 余弦调度
    beta=0.2,                   # 增加 beta
)

# 减少轮数
dpo_config = DPOConfig(num_train_epochs=2)

# 增加 beta
dpo_config = DPOConfig(beta=0.15)

# 早停
from transformers import EarlyStoppingCallback
dpo_trainer = DPOTrainer(
    ...,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=3)]
)