【Qwen2.5+Ollama+PPO+RLHF强化学习训练】通过RLHF方式使用PPO算法，对本地部署的大模型强化学习

概要

强化学习主要使用在SFT后一步，可以减少训练成本并增强大模型的探索能力。本文将介绍从零开始部署模型并RLHF的过程。

整体架构流程

RLHF的PPO算法需要三个模型：奖励模型（评估一个问题的得分）、策略模型（告诉大模型做什么决策更好）、价值模型（告诉大模型做什么事是更有价值的）
奖励模型需要自己训练，而策略模型、价值模型在本文中，笔者通过大模型进行了简单实现。

一、制作数据集

此处笔者准备了一个数据集自动生成器。以芯片尺寸为例。

import random
import pandas as pd
from langchain_ollama import ChatOllama

# 初始化大模型
llm = ChatOllama(model="qwen2.5:7b")

# 定义生成样本数量
NUM_SAMPLES = 1000

# 定义芯片尺寸范围（假设单位为 mm）
MIN_DIM = 1
MAX_DIM = 100

# 定义提示词模板
PROMPT_TEMPLATE = """
你是一个芯片设计顾问，需要根据芯片的尺寸（长、宽、高）推荐调整方向。
输入格式为：[长度, 宽度, 高度]（单位：mm）。
请分析后输出一个 JSON 对象，包含：
- "action": "increase"（增大） / "decrease"（减小） / "no_change"（不动）
- "reason": 简短说明推荐理由（可选）

示例：
输入：[50, 30, 20]
输出：{{"action": "increase", "reason": "散热需求增加"}}

现在请分析以下输入：
输入：{dimensions}
输出：
"""

# 生成数据集
data = []
for i in range(NUM_SAMPLES):
    # 随机生成芯片尺寸
    length = random.randint(MIN_DIM, MAX_DIM)
    width = random.randint(MIN_DIM, MAX_DIM)
    height = random.randint(MIN_DIM, MAX_DIM)

    # 构造输入
    input_dimensions = f"[{length}, {width}, {height}]"

    # 调用大模型生成推荐
    prompt = PROMPT_TEMPLATE.format(dimensions=input_dimensions)
    response = llm.invoke(prompt).content.strip()

    # 解析输出（假设模型返回 JSON 格式）
    try:
        action = "no_change"  # 默认值
        if "increase" in response.lower():
            action = "increase"
        elif "decrease" in response.lower():
            action = "decrease"

        data.append({
            "length": length,
            "width": width,
            "height": height,
            "action": action
        })
    except Exception as e:
        print(f"解析错误（输入：{input_dimensions}）: {e}")
        continue  # 跳过解析失败的样本

# 保存为 CSV 文件
df = pd.DataFrame(data)
df.to_csv("chip_dimension_reward_dataset.csv", index=False)

print(f"已生成 {len(data)} 条样本，保存至 chip_dimension_reward_dataset.csv")

二、训练奖励模型（Reward Model, RM）：

依赖人类标注的偏好数据（如对多个回答的排序或二元比较）。
例如，给定一个问题，人类标注员选择“回答A比回答B更好”，RM 通过学习这些偏好，为每个回答分配一个标量奖励值（如 A 得分 0.8，B 得分 0.3）。
RM 的损失函数通常基于 Bradley-Terry 模型，确保偏好关系与得分一致。

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, accuracy_score
from sklearn.preprocessing import StandardScaler
import joblib

# 1. 加载数据集
data_path = "chip_dimension_reward_dataset.csv"
df = pd.read_csv(data_path)

# 2. 数据预处理
X = df[["length", "width", "height"]]  # 输入特征
y = df["action"]                        # 标签（increase / decrease / no_change）

# 将类别标签编码为整数
y = y.replace({"increase": 0, "decrease": 1, "no_change": 2})

# 可选：标准化输入特征（适用于线性模型）
scaler = StandardScaler()
X = scaler.fit_transform(X)

# 3. 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# 4. 训练模型（这里用随机森林分类器）
model = RandomForestClassifier(
    n_estimators=100,
    max_depth=10,
    random_state=42
)
model.fit(X_train, y_train)

# 5. 模型评估
y_pred = model.predict(X_test)
print("模型评估报告：")
print(classification_report(y_test, y_pred, target_names=["increase", "decrease", "no_change"]))
print(f"准确率: {accuracy_score(y_test, y_pred):.2%}")

# 6. 保存模型和标准化器（便于 PPO 使用）
joblib.dump(model, "reward_model.pkl")
joblib.dump(scaler, "scaler.pkl")

print("模型已训练并保存为 reward_model.pkl 和 scaler.pkl")

训练成功后，使用以下代码测试是否可以完成评测

import joblib
import numpy as np

# 加载模型和标准化器
model = joblib.load("reward_model.pkl")
scaler = joblib.load("scaler.pkl")


# 定义一个函数用于预测芯片尺寸调整方向
def predict_action(length, width, height):
    # 输入特征
    input_data = np.array([[length, width, height]])

    # 标准化输入
    scaled_input = scaler.transform(input_data)

    # 预测动作类别（0: increase, 1: decrease, 2: no_change）
    action_class = model.predict(scaled_input)[0]

    # 可选：返回对应的动作名称
    action_labels = ["increase", "decrease", "no_change"]
    return action_labels[action_class]
print(predict_action(10, 20, 30))
print(predict_action(85, 93, 10))

三、PPO算法

import gym
from gym import spaces
import numpy as np
import pandas as pd
from langchain_ollama import ChatOllama
from stable_baselines3 import PPO
from stable_baselines3.common.policies import ActorCriticPolicy
import torch
import torch.nn as nn

class ChipOptimizationEnv(gym.Env):
    def __init__(self, dataset_path="./dataset/chip_ppo_dataset.csv"):
        super(ChipOptimizationEnv, self).__init__()
        self.llm = ChatOllama(model="qwen2.5:7b")
        self.dataset = pd.read_csv(dataset_path)
        self.current_idx = 0
        self.max_steps = len(self.dataset)

        # 动作空间：0: 增加宽度, 1: 减少宽度, 2: 增加长度, 3: 减少长度
        self.action_space = spaces.Discrete(4)
        # 观察空间：宽度、长度、功耗、性能评分
        self.observation_space = spaces.Box(low=0, high=np.inf, shape=(4,), dtype=np.float32)

        self.current_chip = None
        self.step_count = 0
        self.max_steps_per_episode = 10

    def reset(self):
        self.current_idx = np.random.randint(0, self.max_steps)
        self.current_chip = self.dataset.iloc[self.current_idx].copy()
        self.step_count = 0
        return self._get_observation()

    def _get_observation(self):
        return np.array([
            self.current_chip["width_nm"],
            self.current_chip["length_nm"],
            self.current_chip["power_mw"],
            self.current_chip["performance_score"]
        ], dtype=np.float32)

    def _compute_reward(self, action):
        # 规则奖励：基于性能和功耗
        rule_reward = self.current_chip["performance_score"] / 100.0 - self.current_chip["power_mw"] / 500.0
        # QWEN2.5 奖励：调用模型评估动作
        prompt = (
            f"芯片参数：宽度={self.current_chip['width_nm']:.2f}nm, "
            f"长度={self.current_chip['length_nm']:.2f}nm, "
            f"功耗={self.current_chip['power_mw']:.2f}mW, "
            f"性能评分={self.current_chip['performance_score']:.2f}, "
            f"动作={action}（0: 增加宽度, 1: 减少宽度, 2: 增加长度, 3: 减少长度）。\n"
            f"请评估该动作的优化效果（返回 0 到 1 的评分，1 表示最佳优化）。"
        )
        response = self.llm.invoke(prompt)
        try:
            llm_reward = float(response.content.strip())
        except:
            llm_reward = 0.5  # 默认值
        # 综合奖励
        reward = 0.7 * rule_reward + 0.3 * llm_reward
        if action == self.current_chip["optimal_action"]:
            reward += 0.2  # 额外奖励
        return reward

    def step(self, action):
        self.step_count += 1
        adjustment = 5.0
        # 调整芯片参数
        if action == 0:
            self.current_chip["width_nm"] += adjustment
        elif action == 1:
            self.current_chip["width_nm"] = max(5.0, self.current_chip["width_nm"] - adjustment)
        elif action == 2:
            self.current_chip["length_nm"] += adjustment
        elif action == 3:
            self.current_chip["length_nm"] = max(5.0, self.current_chip["length_nm"] - adjustment)

        # 计算奖励
        reward = self._compute_reward(action)
        done = self.step_count >= self.max_steps_per_episode or self.current_chip["performance_score"] >= 95.0
        info = {"chip_id": self.current_chip["chip_id"]}

        return self._get_observation(), reward, done, info

class CustomPolicy(ActorCriticPolicy):
    def __init__(self, observation_space, action_space, lr_schedule, *args, **kwargs):
        super(CustomPolicy, self).__init__(observation_space, action_space, lr_schedule, *args, **kwargs)
        self.llm = ChatOllama(model="qwen2.5:7b")
        # 自定义策略网络和价值网络
        self.actor = nn.Sequential(
            nn.Linear(observation_space.shape[0], 64),
            nn.ReLU(),
            nn.Linear(64, 64),
            nn.ReLU(),
            nn.Linear(64, action_space.n),
            nn.Softmax(dim=-1)
        )
        self.critic = nn.Sequential(
            nn.Linear(observation_space.shape[0], 64),
            nn.ReLU(),
            nn.Linear(64, 64),
            nn.ReLU(),
            nn.Linear(64, 1)
        )

    def forward(self, obs, deterministic=False):
        obs_tensor = obs if isinstance(obs, torch.Tensor) else torch.tensor(obs, dtype=torch.float32)
        # 策略模型：生成动作分布
        action_probs = self.actor(obs_tensor)
        # 结合 QWEN2.5 建议
        prompt = (
            f"芯片参数：宽度={obs_tensor[0]:.2f}nm, 长度={obs_tensor[1]:.2f}nm, "
            f"功耗={obs_tensor[2]:.2f}mW, 性能评分={obs_tensor[3]:.2f}。\n"
            f"建议调整方向（0: 增加宽度, 1: 减少宽度, 2: 增加长度, 3: 减少长度）："
        )
        response = self.llm.invoke(prompt)
        try:
            action_suggested = int(response.content.strip())
            action_probs[action_suggested] += 0.5  # 增强建议动作的概率
            action_probs = action_probs / action_probs.sum()  # 归一化
        except:
            pass
        # 选择动作
        action = torch.argmax(action_probs, dim=-1) if deterministic else torch.multinomial(action_probs, 1)
        # 价值模型：估计状态值
        value = self.critic(obs_tensor)
        return action, value, torch.log(action_probs)

# 创建环境
env = ChipOptimizationEnv()

# 初始化 PPO 代理
agent = PPO(CustomPolicy, env, verbose=1, learning_rate=3e-4, batch_size=64)

# 训练代理
agent.learn(total_timesteps=10000)

# 保存模型
agent.save("ppo_qwen2.5_chip_model")

训练成功后用以下代码测试

obs = env.reset()
done = False
total_reward = 0
while not done:
    action, _ = agent.predict(obs)
    obs, reward, done, info = env.step(action)
    total_reward += reward
    print(f"Chip ID: {info['chip_id']}, Action: {action}, Reward: {reward:.2f}")
print(f"Total reward: {total_reward:.2f}")

总结

1.PPO，Proximal Policy Optimization，近端策略优化
训练奖励模型（Reward Model, RM）：
依赖人类标注的偏好数据（如对多个回答的排序或二元比较）。
例如，给定一个问题，人类标注员选择“回答A比回答B更好”，RM 通过学习这些偏好，为每个回答分配一个标量奖励值（如 A 得分 0.8，B 得分 0.3）。
RM 的损失函数通常基于 Bradley-Terry 模型，确保偏好关系与得分一致。

强化学习优化策略：
使用 PPO 算法，以 RM 输出的奖励为信号，优化模型策略（Policy）。
PPO 通过限制新旧策略的更新幅度（如 KL 散度约束），避免训练不稳定，逐步提升模型生成高奖励回答的能力。

具体步骤
1.训练reward模型（通常在SFT之后）（通常用能力差不多的大模型）
将输出的LM Head（hidden_size * vocabulary_size）改造为REWARD HEAD（hidden_size * 1）
模型会评估出数据集中的chosen 和 rejected 的值

2.PPO模型
策略模型 价值模型 奖励模型 （其中策略、价值模型都是需要训练的模型，输出是策略模型）
（价值模型用来使模型知道怎么做更好，策略模型让模型知道怎么做决策）
随后比较训练模型和基准模型的KL散度
Reward值 = Score + KL*-0.2
Loss值 采样方法用广义优势法：GAE算法 + V值 使用训练模型、基准模型、重要度采样模型