强化学习与RLHF

import numpy as np
import random
from collections import defaultdict

class QLearningAgent:
    """Q-Learning智能体"""
    
    def __init__(self, state_size, action_size, learning_rate=0.1, 
                 discount_factor=0.99, epsilon=0.1):
        self.state_size = state_size
        self.action_size = action_size
        self.lr = learning_rate
        self.gamma = discount_factor
        self.epsilon = epsilon
        
        # 初始化Q表
        self.q_table = np.zeros((state_size, action_size))
    
    def get_action(self, state):
        """ε-贪婪策略选择动作"""
        if random.random() < self.epsilon:
            # 探索：随机选择动作
            return random.randint(0, self.action_size - 1)
        else:
            # 利用：选择Q值最大的动作
            return np.argmax(self.q_table[state])
    
    def update(self, state, action, reward, next_state, done):
        """更新Q值"""
        if done:
            target = reward
        else:
            target = reward + self.gamma * np.max(self.q_table[next_state])
        
        # Q-Learning更新公式
        self.q_table[state][action] += self.lr * (target - self.q_table[state][action])
    
    def train(self, env, episodes=1000):
        """训练智能体"""
        rewards_per_episode = []
        
        for episode in range(episodes):
            state = env.reset()
            total_reward = 0
            done = False
            
            while not done:
                action = self.get_action(state)
                next_state, reward, done, _ = env.step(action)
                
                self.update(state, action, reward, next_state, done)
                
                state = next_state
                total_reward += reward
            
            rewards_per_episode.append(total_reward)
            
            # 每100个episode打印一次
            if (episode + 1) % 100 == 0:
                avg_reward = np.mean(rewards_per_episode[-100:])
                print(f"Episode {episode + 1}, Average Reward: {avg_reward:.2f}")
        
        return rewards_per_episode

# 示例使用
# env = gym.make('FrozenLake-v1')
# agent = QLearningAgent(state_size=16, action_size=4)
# rewards = agent.train(env, episodes=5000)
print("Q-Learning算法示例代码")

import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical

class PolicyNetwork(nn.Module):
    """策略网络：输入状态，输出动作概率分布"""
    
    def __init__(self, state_size, action_size, hidden_size=64):
        super(PolicyNetwork, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(state_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, action_size),
            nn.Softmax(dim=-1)
        )
    
    def forward(self, state):
        return self.network(state)

class PolicyGradientAgent:
    """策略梯度智能体"""
    
    def __init__(self, state_size, action_size, lr=0.001):
        self.policy = PolicyNetwork(state_size, action_size)
        self.optimizer = optim.Adam(self.policy.parameters(), lr=lr)
        self.gamma = 0.99
        
        # 存储轨迹
        self.states = []
        self.actions = []
        self.rewards = []
    
    def get_action(self, state):
        """根据策略网络选择动作"""
        state_tensor = torch.FloatTensor(state).unsqueeze(0)
        action_probs = self.policy(state_tensor)
        
        # 从概率分布中采样
        dist = Categorical(action_probs)
        action = dist.sample()
        
        # 记录状态和动作
        self.states.append(state_tensor)
        self.actions.append(action)
        
        return action.item()
    
    def store_reward(self, reward):
        """存储奖励"""
        self.rewards.append(reward)
    
    def update(self):
        """使用策略梯度更新网络"""
        # 计算折扣回报
        returns = []
        G = 0
        for reward in reversed(self.rewards):
            G = reward + self.gamma * G
            returns.insert(0, G)
        
        returns = torch.FloatTensor(returns)
        
        # 标准化回报（减少方差）
        if len(returns) > 1:
            returns = (returns - returns.mean()) / (returns.std() + 1e-8)
        
        # 计算策略梯度损失
        policy_loss = []
        for log_prob, G in zip(
            [torch.log(prob) for prob in self.policy(torch.cat(self.states)).gather(1, torch.stack(self.actions).unsqueeze(1)).squeeze()],
            returns
        ):
            policy_loss.append(-log_prob * G)
        
        policy_loss = torch.stack(policy_loss).sum()
        
        # 更新网络
        self.optimizer.zero_grad()
        policy_loss.backward()
        self.optimizer.step()
        
        # 清空轨迹
        self.states = []
        self.actions = []
        self.rewards = []
        
        return policy_loss.item()

print("策略梯度算法示例代码")

import torch
import torch.nn as nn
from transformers import AutoModelForSequenceClassification, AutoTokenizer

class RewardModel(nn.Module):
    """奖励模型：基于预训练语言模型"""
    
    def __init__(self, model_name="gpt2", hidden_size=768):
        super(RewardModel, self).__init__()
        self.reward_model = AutoModelForSequenceClassification.from_pretrained(
            model_name, num_labels=1
        )
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
    
    def forward(self, input_ids, attention_mask):
        """前向传播，返回奖励分数"""
        outputs = self.reward_model(input_ids, attention_mask=attention_mask)
        return outputs.logits.squeeze(-1)

def reward_loss(reward_chosen, reward_rejected):
    """Bradley-Terry偏好损失"""
    # 希望 chosen 的奖励高于 rejected
    loss = -torch.log(torch.sigmoid(reward_chosen - reward_rejected))
    return loss.mean()

def train_reward_model(model, train_dataloader, optimizer, epochs=3):
    """训练奖励模型"""
    model.train()
    
    for epoch in range(epochs):
        total_loss = 0
        
        for batch in train_dataloader:
            # 编码 chosen 和 rejected 回答
            chosen_input = model.tokenizer(
                batch['chosen'], 
                padding=True, 
                truncation=True, 
                return_tensors='pt'
            )
            rejected_input = model.tokenizer(
                batch['rejected'], 
                padding=True, 
                truncation=True, 
                return_tensors='pt'
            )
            
            # 计算奖励
            reward_chosen = model(**chosen_input)
            reward_rejected = model(**rejected_input)
            
            # 计算损失
            loss = reward_loss(reward_chosen, reward_rejected)
            
            # 更新模型
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        avg_loss = total_loss / len(train_dataloader)
        print(f"Epoch {epoch + 1}, Average Loss: {avg_loss:.4f}")

print("奖励模型训练示例代码")

import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import numpy as np

class PPOAgent:
    """PPO智能体：用于大语言模型的RLHF训练"""
    
    def __init__(self, policy_model, ref_model, reward_model, 
                 lr=1e-5, clip_range=0.2, gamma=0.99):
        self.policy_model = policy_model  # 被优化的策略模型
        self.ref_model = ref_model        # 参考模型（SFT模型）
        self.reward_model = reward_model  # 奖励模型
        
        self.optimizer = optim.Adam(self.policy_model.parameters(), lr=lr)
        self.clip_range = clip_range
        self.gamma = gamma
        
        self.kl_coef = 0.1  # KL散度惩罚系数
    
    def compute_kl_penalty(self, policy_logits, ref_logits):
        """计算KL散度惩罚"""
        policy_dist = torch.distributions.Categorical(logits=policy_logits)
        ref_dist = torch.distributions.Categorical(logits=ref_logits)
        
        # KL(policy || ref)
        kl = torch.distributions.kl_divergence(policy_dist, ref_dist)
        return kl
    
    def ppo_loss(self, old_log_probs, new_log_probs, advantages, clip_range):
        """PPO损失函数"""
        # 比率：新策略 / 旧策略
        ratio = torch.exp(new_log_probs - old_log_probs)
        
        # 裁剪的目标函数
        clipped_ratio = torch.clamp(ratio, 1 - clip_range, 1 + clip_range)
        
        # 取两者中的最小值
        surrogate1 = ratio * advantages
        surrogate2 = clipped_ratio * advantages
        
        return -torch.min(surrogate1, surrogate2).mean()
    
    def compute_advantages(self, rewards, values, gamma=0.99, lam=0.95):
        """计算优势函数（GAE）"""
        advantages = []
        gae = 0
        
        for reward, value in zip(reversed(rewards), reversed(values)):
            next_value = values[values.index(value) + 1] if values.index(value) < len(values) - 1 else 0
            delta = reward + gamma * next_value - value
            gae = delta + gamma * lam * gae
            advantages.insert(0, gae)
        
        advantages = torch.FloatTensor(advantages)
        
        # 标准化
        if len(advantages) > 1:
            advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        
        return advantages
    
    def update(self, rollout):
        """PPO更新"""
        states = rollout['states']
        actions = rollout['actions']
        old_log_probs = rollout['log_probs']
        rewards = rollout['rewards']
        
        # 计算当前策略的log概率
        new_log_probs = []
        policy_logits_list = []
        ref_logits_list = []
        
        for state in states:
            # 当前策略
            state_tensor = torch.FloatTensor(state).unsqueeze(0)
            policy_output = self.policy_model(state_tensor)
            policy_logits = policy_output['logits']
            policy_dist = torch.distributions.Categorical(logits=policy_logits)
            new_log_prob = policy_dist.log_prob(torch.LongTensor([actions[states.index(state)]]))
            new_log_probs.append(new_log_prob)
            policy_logits_list.append(policy_logits)
            
            # 参考模型
            with torch.no_grad():
                ref_output = self.ref_model(state_tensor)
                ref_logits_list.append(ref_output['logits'])
        
        new_log_probs = torch.stack(new_log_probs)
        old_log_probs = torch.FloatTensor(old_log_probs)
        
        # 计算优势函数
        values = torch.zeros(len(rewards))  # 简化：假设value=0
        advantages = self.compute_advantages(rewards, values)
        
        # PPO损失
        ppo_loss = self.ppo_loss(old_log_probs, new_log_probs, advantages, self.clip_range)
        
        # KL惩罚
        kl_penalties = []
        for policy_logits, ref_logits in zip(policy_logits_list, ref_logits_list):
            kl_penalty = self.compute_kl_penalty(policy_logits, ref_logits)
            kl_penalties.append(kl_penalty)
        kl_penalty = torch.stack(kl_penalties).mean()
        
        # 总损失
        total_loss = ppo_loss + self.kl_coef * kl_penalty
        
        # 更新模型
        self.optimizer.zero_grad()
        total_loss.backward()
        self.optimizer.step()
        
        return ppo_loss.item(), kl_penalty.item()

print("PPO算法实现示例")

import torch
import torch.nn.functional as F

def dpo_loss(policy_chosen_logps, policy_rejected_logps,
             reference_chosen_logps, reference_rejected_logps,
             beta=0.1):
    """
    DPO损失函数
    
    Args:
        policy_chosen_logps: 策略模型对chosen回答的log概率
        policy_rejected_logps: 策略模型对rejected回答的log概率
        reference_chosen_logps: 参考模型对chosen回答的log概率
        reference_rejected_logps: 参考模型对rejected回答的log概率
        beta: 温度参数，控制与参考模型的偏离程度
    
    Returns:
        DPO损失
    """
    # 计算log比率
    chosen_log_ratio = policy_chosen_logps - reference_chosen_logps
    rejected_log_ratio = policy_rejected_logps - reference_rejected_logps
    
    # DPO损失
    logits = beta * (chosen_log_ratio - rejected_log_ratio)
    loss = -F.logsigmoid(logits).mean()
    
    return loss

# 计算log概率的辅助函数
def compute_log_probs(model, input_ids, attention_mask):
    """计算模型对输入的log概率"""
    outputs = model(input_ids, attention_mask=attention_mask)
    logits = outputs.logits
    
    # 使用log_softmax计算log概率
    log_probs = F.log_softmax(logits, dim=-1)
    
    # 取token级别的log概率并求和
    token_log_probs = torch.gather(
        log_probs[:, :-1], 
        dim=2, 
        index=input_ids[:, 1:].unsqueeze(2)
    ).squeeze(2)
    
    return token_log_probs.sum(dim=-1)

print("DPO损失函数示例")

# 安装TRL及其依赖
pip install trl
pip install transformers
pip install datasets
pip install peft  # 用于LoRA

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from trl import SFTTrainer
from datasets import load_dataset

# 加载模型和分词器
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# 加载数据集
dataset = load_dataset("imdb", split="train[:1000]")

def preprocess_function(examples):
    """预处理数据集"""
    return tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=512
    )

tokenized_dataset = dataset.map(preprocess_function, batched=True)

# 训练参数
training_args = TrainingArguments(
    output_dir="./sft_model",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    learning_rate=2e-5,
    warmup_steps=100,
    logging_steps=10,
    save_strategy="epoch",
)

# SFT训练器
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    tokenizer=tokenizer,
    max_seq_length=512,
)

# 开始训练
trainer.train()
print("SFT训练完成！")

from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead
from transformers import AutoTokenizer
import torch

# 1. 加载SFT模型（作为初始策略）
model_name = "gpt2"
model = AutoModelForCausalLMWithValueHead.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# 2. 加载参考模型（SFT模型的副本）
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(model_name)

# 3. PPO配置
ppo_config = PPOConfig(
    learning_rate=1e-5,
    batch_size=16,
    mini_batch_size=4,
    ppo_epochs=4,
    kl_penalty="kl",  # 或 "abs", "full"
    init_kl_coef=0.2,
    target_kl=6.0,
)

# 4. 创建PPO训练器
ppo_trainer = PPOTrainer(
    config=ppo_config,
    model=model,
    ref_model=ref_model,
    tokenizer=tokenizer,
)

# 5. 加载奖励模型（示例：使用简单的启发式奖励）
def reward_function(response):
    """简单的奖励函数（实际应使用训练好的奖励模型）"""
    # 示例：奖励包含特定关键词的回答
    good_keywords = ["helpful", "accurate", "clear"]
    reward = sum(1 for word in good_keywords if word in response.lower())
    return reward

# 6. 准备训练数据
prompts = [
    "What is machine learning?",
    "Explain neural networks",
    "What is reinforcement learning?",
]

# 7. RLHF训练循环
for epoch in range(3):
    print(f"\n=== Epoch {epoch + 1} ===")
    
    for prompt in prompts:
        # 编码prompt
        query_tensors = tokenizer.encode(prompt, return_tensors="pt")[0]
        
        # 生成回答
        response_tensors = ppo_trainer.generate(
            query_tensors.unsqueeze(0),
            max_new_tokens=128,
            do_sample=True,
            top_k=50,
            top_p=0.95,
        )
        
        # 解码回答
        response = tokenizer.decode(response_tensors[0], skip_special_tokens=True)
        
        # 计算奖励
        reward = reward_function(response)
        reward_tensors = [torch.tensor([reward])]
        
        # PPO更新
        stats = ppo_trainer.step(
            [query_tensors],
            [response_tensors[0]],
            reward_tensors
        )
        
        print(f"Prompt: {prompt[:30]}...")
        print(f"Reward: {reward:.2f}")
        print(f"KL: {stats['ppo/kl']:.2f}")
        print(f"Entropy: {stats['ppo/entropy']:.2f}")

print("\nRLHF训练完成！")