下面给出一个更完整、带有多样化交互策略的强化学习示例。相较于之前只对最终输出给一个标量的简单做法，本示例将展示如何在训练过程中：

1. 将「环境」(Environment) 与「智能体」(Agent) 的概念更清晰地拆分——「环境」向智能体提供状态、给出奖励，「智能体」基于策略选择动作(也就是生成文本的过程)，并在一批轨迹结束时，基于优势函数(Advantage function) 对策略进行更新。
2. 通过「动作」(Actions) 的概念，来表示每一步token的选择（或一段回复的生成）。
3. 计算多元组合的奖励(Rewards)，比如使用之前的 dynamic_reward_func 做整体的质量评分，再结合一些额外交互(如人工反馈、交互式打分、或对环境中其他信息的采样) 来组成总奖励。
4. 计算「优势」(Advantage) 以及「价值函数」(Value function) 的误差，通过PPO算法(或兼容TRL的GRPO)进行策略更新。

注意：此示例在结构上仍然沿用之前的思路（TRL库中的PPO/GRPO流程），只是在逻辑上补充了更多「环境-智能体」的拆分、rollout缓存、优势计算的过程示范，以帮助你理解和扩展。实际使用中可根据自己的需求精简或修改。

下面的代码是一个完整示例，从头到尾包括：

• 数据加载与预处理
• 模型与Tokenizer构建
• 环境(Environment) 和 智能体(Agent) 的定义
• 回合(episode) 采样、动作生成、奖励计算
• 优势函数(Advantage) 估计
• 利用PPO/GRPO进行策略网络更新
• 最后对模型进行推断和评估

请注意，由于此示例会同时包含前面所涉及的各项奖励函数以及额外的「环境/智能体/优势函数」逻辑，会比较长。但为了**“不要有任何的省略”**，我会将其完整呈现：

python
# -*- coding: utf-8 -*-
import torch
import re
import random
import langid
from typing import List, Dict, Any
import math

import nltk
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
from nltk import edit_distance

# 如果你没有安装 nltk 或 langid，需要先安装:
#   pip install nltk langid

# 如果你没有安装 datasets，需要先安装:
#   pip install datasets
from datasets import load_dataset, Dataset, concatenate_datasets

# 如果你没有安装 trl，需要先安装（确保版本>=0.4）:
#   pip install trl

from trl import PPOConfig, PPOTrainer, AutoModelForCausalLMWithValueHead, AutoTokenizer
# 或者如果你在使用 GRPO:
# from trl import GRPOConfig, GRPOTrainer

# 如果你有 wandb，需要的话可以导入并在训练中使用:
import wandb
import os

# ==============
# (可选)仅用于指定GPU。如果你的环境无需此操作可注释
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

#######################
#   自定义一些配置常量
#######################
SUPPORTED_LANGUAGES = {
    "en_to_zh": ("英语", "中文"),
    "zh_to_en": ("中文", "英语"),
    "jp_to_zh": ("日语", "中文"),
    "zh_to_azh": ("中文", "中文（文言文）"),
    "en_to_azh": ("英语", "中文（文言文）"),
    "de_to_zh": ("德语", "中文"),
    "italy_to_zh": ("意大利语", "中文"),
    "rus_to_zh": ("俄语", "中文"),
    "fr_to_zh": ("法语", "中文"),
    "mix": ("中文", "multilingual")
}

SYSTEM_PROMPT = """\
You are a versatile assistant. Follow these rules:  

1. **Always** respond in this exact XML format:  
   <reasoning>
   [Your reasoning here]
   </reasoning>
   <answer>
   [Final answer or translation here]
   </answer>
2. If given lyrics, translate them from {} to {} only, without extra explanations.  
3. If given a math question, solve it and provide the result in the same XML format.
"""

XML_COT_FORMAT = """\
<reasoning>
{reasoning}
</reasoning>
<answer>
{answer}
</answer>
"""

#######################
#       数据准备函数
#######################
def get_lyric_datasets(path: str, train_test_ratio=0.8) -> dict:
    """
    加载一个歌词翻译数据集 JSON，并拆分成 train / test (例如80% / 20%)。
    每条数据应包含 'type', 'lyric', 'target_lyric' 这些字段。
    本示例在map_fn里插入了对应的系统提示 + user内容，并返回 (prompt, answer, task)
    """
    data = Dataset.from_json(path)

    def map_fn(x):
        lang_src = SUPPORTED_LANGUAGES[x['type']][0]
        lang_tgt = SUPPORTED_LANGUAGES[x['type']][1]
        # 将正确的源语言->目标语言插入 system prompt:
        system_plus = SYSTEM_PROMPT.format(lang_src, lang_tgt)

        return {
            'prompt': [
                {'role': 'system', 'content': system_plus},
                {'role': 'user', 'content': x['lyric']}
            ],
            'answer': x['target_lyric'],
            'task': "lyrics"
        }

    data = data.map(map_fn)
    return data.train_test_split(train_size=train_test_ratio, seed=42)

def extract_hash_answer(text: str):
    """
    对于 GSM8K数据常见格式：在解析过程中提取 '####' 后的答案部分。
    """
    if "####" not in text:
        return None
    return text.split("####")[1].strip()

def get_gsm8k_questions(split="train", sample_size=3000) -> Dataset:
    """
    加载 GSM8K 并统一处理 prompt、answer。仅供示例，需根据你真实路径或方式改写。
    """
    # 此处假设你有若干 gsm8k的parquet文件: ../data_pack/openai/gsm8k/train-*.parquet
    # 纯演示
    data = Dataset.from_parquet(f'../data_pack/openai/gsm8k/{split}-*.parquet')
    data = data.select(range(min(sample_size, len(data))))

    def map_fn(x):
        user_prompt = x['question']
        final_answer = extract_hash_answer(x['answer'])
        return {
            'prompt': [
                {'role': 'system', 'content': SYSTEM_PROMPT.format("N/A", "N/A")},
                {'role': 'user', 'content': user_prompt}
            ],
            'answer': final_answer if final_answer else "",
            'task': "math"
        }

    data = data.map(map_fn)
    return data

#######################
#   文本抽取相关小工具
#######################
def extract_xml_answer(text: str) -> str:
    """
    从模型生成的文本中提取<answer>...</answer>之间的内容。
    """
    pattern = r"<answer>\s*(.*?)\s*</answer>"
    match = re.search(pattern, text, re.DOTALL | re.IGNORECASE)
    if match:
        return match.group(1).strip()
    return ""

def detect_language(text: str) -> str:
    """
    langid 做简单的语言检测, 返回 'en'、'zh'、'fr'等ISO语言代码
    """
    return langid.classify(text)[0]


#######################
#   奖励函数示例
#######################
def compute_length_acc(
    preds: List[str],
    refs: List[str],
    tokenizer,
    max_tolerance: float = 0.5
) -> List[float]:
    """
    返回每个pred/ref的长度相似度奖励。
    仅作为示例，可根据需求自行调整。
    """
    rewards = []
    for pred, ref in zip(preds, refs):
        pred_tokens = tokenizer.tokenize(pred)
        ref_tokens = tokenizer.tokenize(ref)
        if len(ref_tokens) == 0:
            rewards.append(0.0)
            continue
        length_ratio = abs(len(pred_tokens) - len(ref_tokens)) / len(ref_tokens)

        if length_ratio <= 0.1:
            score = 1.0
        elif length_ratio <= 0.2:
            score = 0.7
        elif length_ratio <= 0.3:
            score = 0.5
        elif length_ratio <= 0.4:
            score = 0.3
        else:
            score = 0.0

        rewards.append(score)
    return rewards

def compute_bleu(
    preds: List[str],
    refs: List[str],
    tokenizer
) -> List[float]:
    """
    逐个pair计算BLEU(4-gram平均)
    """
    smoothie = SmoothingFunction().method1
    weights = (0.25, 0.25, 0.25, 0.25)
    scores = []
    for pred, ref in zip(preds, refs):
        pred_tokens = tokenizer.tokenize(pred)
        ref_tokens = tokenizer.tokenize(ref)
        if not pred_tokens or not ref_tokens:
            scores.append(0.0)
            continue
        bleu = sentence_bleu(
            [ref_tokens],
            pred_tokens,
            weights=weights,
            smoothing_function=smoothie
        )
        scores.append(bleu)
    return scores

def compute_ter(
    preds: List[str],
    refs: List[str],
    tokenizer
) -> List[float]:
    """
    计算 Translation Edit Rate = 编辑距离 / 参考词数 * 100
    """
    ter_scores = []
    for pred, ref in zip(preds, refs):
        pred_tokens = tokenizer.tokenize(pred)
        ref_tokens = tokenizer.tokenize(ref)
        if len(ref_tokens) == 0:
            if len(pred_tokens) > 0:
                ter_scores.append(100.0)
            else:
                ter_scores.append(0.0)
            continue
        dist = edit_distance(pred_tokens, ref_tokens)
        ter = (dist / len(ref_tokens)) * 100
        ter_scores.append(ter)
    return ter_scores


# ===== 一些基础奖励函数 =====
def strict_format_reward_func(completions, **kwargs) -> List[float]:
    """
    如果文本完整匹配 <reasoning>...</reasoning><answer>...</answer> 格式 => +1，否则0
    """
    pattern = r"<reasoning>[\s\S]+</reasoning>\s*<answer>[\s\S]+</answer>"
    responses = [completion[0]["content"] for completion in completions]
    scores = []
    for r in responses:
        if re.search(pattern, r):
            scores.append(1.0)
        else:
            scores.append(0.0)
    return scores

def soft_format_reward_func(completions, **kwargs) -> List[float]:
    """
    包含最基本XML标签 <reasoning>...</reasoning> <answer>...</answer> 就给0.5，否则0
    """
    pattern = r"<reasoning>.*?</reasoning>\s*<answer>.*?</answer>"
    responses = [completion[0]["content"] for completion in completions]
    matches = [re.search(pattern, r, flags=re.DOTALL) for r in responses]
    return [0.5 if m else 0.0 for m in matches]

def xmlcount_reward_func(completions, **kwargs) -> List[float]:
    """
    根据标签出现次数给分，演示用。
    """
    def count_xml(text) -> float:
        count = 0.0
        if text.count("<reasoning>\n") == 1:
            count += 0.125
        if text.count("\n</reasoning>\n") == 1:
            count += 0.125
        if text.count("\n<answer>\n") == 1:
            count += 0.125
        if text.count("\n</answer>") == 1:
            count += 0.125
        return count

    responses = [c[0]["content"] for c in completions]
    return [count_xml(c) for c in responses]

# ===== 数学任务的奖励函数 =====
def correctness_reward_func(prompts, completions, answers, **kwargs) -> List[float]:
    """
    如果提取到的 <answer> 和正确答案一致 => +2, 否则0
    """
    responses = [c[0]['content'] for c in completions]
    extracted_responses = [extract_xml_answer(r) for r in responses]
    return [2.0 if r.strip() == a.strip() else 0.0 for r, a in zip(extracted_responses, answers)]

def int_reward_func(completions, **kwargs) -> List[float]:
    """
    如果 <answer> 是纯数字(可带符号)，奖励+0.5，否则0
    """
    responses = [c[0]['content'] for c in completions]
    extracted_responses = [extract_xml_answer(r) for r in responses]
    scores = []
    for r in extracted_responses:
        if re.match(r"^-?\d+$", r.strip()):
            scores.append(0.5)
        else:
            scores.append(0.0)
    return scores

# ===== 歌词翻译任务的奖励函数 =====
def length_acc_reward_func(prompts, completions, answers, tokenizer=None, **kwargs) -> List[float]:
    responses = [completion[0]['content'] for completion in completions]
    extracted = [extract_xml_answer(r) for r in responses]
    return compute_length_acc(preds=extracted, refs=answers, tokenizer=tokenizer)

def bleu_reward_func(prompts, completions, answers, tokenizer=None, **kwargs) -> List[float]:
    responses = [c[0]["content"] for c in completions]
    extracted = [extract_xml_answer(r) for r in responses]
    bleu_scores = compute_bleu(preds=extracted, refs=answers, tokenizer=tokenizer)
    rewards = []
    for s in bleu_scores:
        if s >= 0.9:
            rewards.append(5.0)
        elif s >= 0.8:
            rewards.append(4.5)
        elif s >= 0.7:
            rewards.append(4.0)
        elif s >= 0.6:
            rewards.append(3.5)
        elif s >= 0.5:
            rewards.append(2.5)
        elif s >= 0.4:
            rewards.append(2.0)
        elif s >= 0.3:
            rewards.append(1.5)
        elif s >= 0.2:
            rewards.append(1.0)
        elif s >= 0.1:
            rewards.append(0.5)
        else:
            rewards.append(0.0)
    return rewards

def ter_reward_func(prompts, completions, answers, tokenizer=None, **kwargs) -> List[float]:
    responses = [c[0]["content"] for c in completions]
    extracted = [extract_xml_answer(r) for r in responses]
    ter_scores = compute_ter(preds=extracted, refs=answers, tokenizer=tokenizer)
    rewards = []
    for t in ter_scores:
        if t >= 80:
            rewards.append(0.0)
        elif t >= 60:
            rewards.append(0.5)
        elif t >= 40:
            rewards.append(1.0)
        elif t >= 20:
            rewards.append(1.5)
        else:
            rewards.append(2.0)
    return rewards

def language_recognition(prompts, completions, answers, **kwargs) -> List[float]:
    """
    如果生成的文本语言与参考答案语言一致 => +1，否则0
    """
    responses = [c[0]["content"] for c in completions]
    extracted = [extract_xml_answer(r) for r in responses]
    rewards = []
    for pred, ref in zip(extracted, answers):
        if not pred.strip():
            rewards.append(0.0)
            continue
        pred_lang = detect_language(pred)
        ref_lang = detect_language(ref)
        rewards.append(1.0 if pred_lang == ref_lang else 0.0)
    return rewards

def reduce_repetition_reward_func(prompts, completions, answers, **kwargs) -> List[float]:
    """
    简单示例：歌词中若重复行过多 => 惩罚；重复行少 => 奖励
    """
    responses = [c[0]["content"] for c in completions]
    rewards = []
    for resp in responses:
        answer_text = extract_xml_answer(resp)
        lines = [l.strip() for l in answer_text.split("\n") if l.strip()]
        if len(lines) <= 1:
            rewards.append(0.0)
            continue
        unique_lines = len(set(lines))
        ratio = unique_lines / len(lines)
        reward_val = 1.0 if ratio >= 0.8 else 0.0
        rewards.append(reward_val)
    return rewards

# ===== 其他自定义函数 =====
def fear_of_unknown_reward_func(prompts, completions, answers, **kwargs) -> List[float]:
    """
    如果回答中出现“不确定”、“不知道”、“???”等 => 负奖励
    """
    responses = [c[0]["content"] for c in completions]
    triggers = ["i don't know", "i do not know", "not sure", "???", "no idea", "不确定", "不知道"]
    scores = []
    for resp in responses:
        lower_resp = resp.lower()
        penalty = 0.0
        for t in triggers:
            if t in lower_resp:
                penalty -= 1.0  # 每出现一次触发词，-1分
        scores.append(penalty)
    return scores

def contradiction_penalty_func(prompts, completions, answers, **kwargs) -> List[float]:
    """
    若同一回答中出现明显矛盾(如"yes"和"no") => 惩罚
    """
    contradictory_pairs = [
        ("yes", "no"),
        ("true", "false"),
        ("can", "cannot")
    ]
    responses = [c[0]["content"].lower() for c in completions]
    scores = []
    for r in responses:
        penalty = 0.0
        for pair in contradictory_pairs:
            if pair[0] in r and pair[1] in r:
                penalty -= 2.0
        scores.append(penalty)
    return scores

def curiosity_question_reward_func(prompts, completions, answers, **kwargs) -> List[float]:
    """
    若 <reasoning> 中出现对自身问题的好奇(带'?') => 正向奖励
    """
    responses = [c[0]["content"] for c in completions]
    pattern = r"<reasoning>(.*?)</reasoning>"
    scores = []
    for resp in responses:
        match = re.search(pattern, resp, flags=re.DOTALL | re.IGNORECASE)
        if not match:
            scores.append(0.0)
            continue
        reasoning_text = match.group(1)
        question_count = reasoning_text.count('?')
        reward = min(question_count * 0.2, 1.0)
        scores.append(reward)
    return scores

def novelty_reward_func(prompts, completions, answers, tokenizer=None, **kwargs) -> List[float]:
    """
    统计与用户prompt重叠的词，用于鼓励一定程度新意
    """
    user_texts = []
    for p in prompts:
        user_texts.append(p[-1]['content'])

    responses = [c[0]["content"] for c in completions]
    scores = []
    for user_t, resp in zip(user_texts, responses):
        user_tokens = set(tokenizer.tokenize(user_t.lower()))
        model_answer_text = extract_xml_answer(resp).lower()
        model_tokens = tokenizer.tokenize(model_answer_text)
        if not model_tokens:
            scores.append(0.0)
            continue
        new_words = [w for w in model_tokens if w not in user_tokens]
        ratio = len(new_words) / len(model_tokens)
        if ratio > 0.8:
            scores.append(1.0)
        elif ratio > 0.5:
            scores.append(0.5)
        elif ratio > 0.2:
            scores.append(0.2)
        else:
            scores.append(0.0)
    return scores


# ================
#   重点：动态奖励
# ================
def dynamic_reward_func(
    prompts,
    completions,
    answers,
    tokenizer=None,
    batch_metadata=None,  # 如果在Trainer中可以传递额外元数据
    **kwargs
) -> List[float]:
    """
    以任务类型、恐惧检测结果、好奇程度等为依据，动态决定如何组合子奖励函数。
    """

    batch_size = len(completions)
    final_rewards = [0.0] * batch_size

    # 我们先推断哪些是math任务，哪些是lyrics任务
    tasks = []
    if batch_metadata is not None:
        tasks = [meta.get('task', None) for meta in batch_metadata]
    else:
        # 不可用就简单从 prompts 猜测
        for i in range(batch_size):
            user_text = prompts[i][-1]['content'].lower()
            if "math" in user_text or "question" in user_text:
                tasks.append("math")
            else:
                tasks.append("lyrics")

    math_indices = [i for i,t in enumerate(tasks) if t == "math"]
    lyrics_indices = [i for i,t in enumerate(tasks) if t == "lyrics"]

    # 对 math 的样本计算
    if math_indices:
        sub_prompts = [prompts[i] for i in math_indices]
        sub_completions = [completions[i] for i in math_indices]
        sub_answers = [answers[i] for i in math_indices]

        c_scores = correctness_reward_func(sub_prompts, sub_completions, sub_answers)
        i_scores = int_reward_func(sub_completions)
        s_scores = strict_format_reward_func(sub_completions)
        so_scores = soft_format_reward_func(sub_completions)

        sub_final = [
            c + i + s + so
            for c, i, s, so in zip(c_scores, i_scores, s_scores, so_scores)
        ]
        for idx, score_val in zip(math_indices, sub_final):
            final_rewards[idx] += score_val

    # 对 lyrics 的样本计算
    if lyrics_indices:
        sub_prompts = [prompts[i] for i in lyrics_indices]
        sub_completions = [completions[i] for i in lyrics_indices]
        sub_answers = [answers[i] for i in lyrics_indices]

        if tokenizer is None:
            raise ValueError("Tokenizer is required for lyrics reward functions.")

        x_scores = xmlcount_reward_func(sub_completions)
        so_scores = soft_format_reward_func(sub_completions)
        st_scores = strict_format_reward_func(sub_completions)
        l_scores = length_acc_reward_func(sub_prompts, sub_completions, sub_answers, tokenizer=tokenizer)
        b_scores = bleu_reward_func(sub_prompts, sub_completions, sub_answers, tokenizer=tokenizer)
        t_scores = ter_reward_func(sub_prompts, sub_completions, sub_answers, tokenizer=tokenizer)
        lang_scores = language_recognition(sub_prompts, sub_completions, sub_answers)
        rep_scores = reduce_repetition_reward_func(sub_prompts, sub_completions, sub_answers)

        sub_final = [
            x + so + st + l + b + t + la + r
            for x, so, st, l, b, t, la, r in zip(
                x_scores, so_scores, st_scores,
                l_scores, b_scores, t_scores,
                lang_scores, rep_scores
            )
        ]
        for idx, score_val in zip(lyrics_indices, sub_final):
            final_rewards[idx] += score_val

    # 对所有样本添加通用奖励或惩罚
    fear_scores = fear_of_unknown_reward_func(prompts, completions, answers)
    contra_scores = contradiction_penalty_func(prompts, completions, answers)
    curious_scores = curiosity_question_reward_func(prompts, completions, answers)
    novelty_scores = [0.0]*batch_size
    if tokenizer is not None:
        novelty_scores = novelty_reward_func(prompts, completions, answers, tokenizer=tokenizer)

    for i in range(batch_size):
        final_rewards[i] += fear_scores[i]
        final_rewards[i] += contra_scores[i]
        final_rewards[i] += curious_scores[i]
        final_rewards[i] += novelty_scores[i]

    return final_rewards


##############################
#  环境(Env) 与 智能体(Agent)
##############################
class TextRLTaskEnvironment:
    """
    一个简化的示例环境。它会：
      1. 提供一个(batch)的 "state" (即 prompt) 给智能体
      2. 智能体生成动作(即输出文本)
      3. 环境将 "完成的对话" + "参考答案" 传入奖励函数，计算reward
      4. 返回 (obs, reward, done, info)
    这里我们一次生成完整回复(而非token-level循环)，以示例如何与上面的 reward_func 集成。
    更细粒度的token-level循环也可实现，但需要在TRL中自定义更多组件。
    """

    def __init__(self, dataset, tokenizer, reward_func):
        self.dataset = dataset
        self.tokenizer = tokenizer
        self.reward_func = reward_func

    def reset(self, indices):
        """
        给定一批索引，从dataset里取相应的数据，并视为state
        这里只返回 'prompt' 与 'answer' 做参考
        """
        batch = [self.dataset[i] for i in indices]
        # obs可以是prompt, 以及其它元数据
        obs = [ex["prompt"] for ex in batch]
        references = [ex["answer"] for ex in batch]
        # 也可以把task等信息当作metadata
        meta = batch
        return obs, references, meta

    def step(self, prompts, completions, references, meta):
        """
        计算奖励
        prompts: List[List[Dict]]   # role-based prompt
        completions: List[List[Dict]]  # role=assistant, content=xxx
        references: List[str]
        meta: List[dict]   # 里面包含task等
        """
        # reward_func需要的入参有 prompts, completions, answers (就是references)
        # 还可以传 batch_metadata=meta
        rewards = self.reward_func(prompts, completions, references, tokenizer=self.tokenizer, batch_metadata=meta)
        # done：这里一次生成就结束
        done = [True]*len(prompts)
        # info可返回一些诊断信息
        info = {}
        return rewards, done, info


class TextRLLanguageAgent:
    """
    简化示例：使用一个Policy+Value网络(AutoModelForCausalLMWithValueHead)进行文本生成。
    也可以对应到TRLError/PPOTrainer中处理。
    """

    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
        self.model.eval()

    def act(self, prompts, max_new_tokens=128):
        """
        给定一批prompt，使用当前策略生成回复。
        这里演示一次性生成文本。更精细的 token-level act 也可实现。
        """
        # TRL中的PPOTrainer会对生成进行封装，这里只是演示
        completions = []
        device = next(self.model.parameters()).device

        for p in prompts:
            # 将 role-based prompt 拼到一起
            text_input = ""
            for pm in p:
                if pm['role'] == 'system':
                    text_input += f"[system] {pm['content']}\n"
                elif pm['role'] == 'user':
                    text_input += f"[user] {pm['content']}\n"
                else:
                    text_input += f"[{pm['role']}] {pm['content']}\n"

            inputs = self.tokenizer(text_input, return_tensors="pt").to(device)
            with torch.no_grad():
                output_ids = self.model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                    do_sample=True,
                    top_k=50,
                    top_p=0.9,
                    pad_token_id=self.tokenizer.eos_token_id
                )
            gen_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
            # 这里简单取生成文本中 prompt 之后的新内容：
            completion_str = gen_text[len(text_input):].strip()
            completions.append([{'role': 'assistant', 'content': completion_str}])

        return completions


##################################
#   示范一个 PPO/GRPO 训练流程
##################################
def evaluate_model(trainer, test_dataset, tokenizer):
    """
    在测试集上做推断和计算一些指标 (例如 BLEU, TER, length_acc)
    """
    print("\nBegin Evaluation on Test Dataset...")

    model = trainer.model
    model.eval()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    predictions = []
    references = []

    for example in test_dataset:
        prompt = example["prompt"]
        true_answer = example["answer"]
        # 将role-based prompt拼合
        text_input = ""
        for pm in prompt:
            text_input += f"{pm['role']}:{pm['content']}\n"
        inputs = tokenizer(text_input, return_tensors="pt").to(device)

        with torch.no_grad():
            generated_ids = model.generate(**inputs, max_new_tokens=256)
        generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)

        predictions.append(generated_text)
        references.append(true_answer)

    # 提取 <answer>
    pred_answers = [extract_xml_answer(pred) for pred in predictions]

    # 计算评估指标
    bleu_scores = compute_bleu(pred_answers, references, tokenizer=tokenizer)
    ter_scores = compute_ter(pred_answers, references, tokenizer=tokenizer)
    length_accs = compute_length_acc(pred_answers, references, tokenizer=tokenizer)

    # 记录结果
    results = {
        "BLEU": bleu_scores,
        "TER": ter_scores,
        "LENGTH_ACCURACY": length_accs,
    }
    for metric, scores in results.items():
        avg_score = sum(scores) / len(scores) if scores else 0.0
        print(f"[Test] {metric}: Avg {avg_score:.4f}")
    print("Test evaluation completed.")


def main():
    # 1. 准备模型 & tokenizer
    #    这里示例使用HF任意的CausalLM模型; 也可以使用Qwen等，只要能兼容PPOTrainer
    model_name = "gpt2"  # 仅作演示
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # 在trl库里，PPO需要一个带价值头(ValueHead)的模型
    #   也可以使用 GRPO:  from trl import GRPOTrainer, GRPOConfig, ...
    model = AutoModelForCausalLMWithValueHead.from_pretrained(model_name)

    # 2. 准备数据：math 和 lyrics 的例子(自行替换路径)
    print("Loading GSM8K dataset for math tasks...")
    gsm8k_dataset = get_gsm8k_questions(split="train", sample_size=3000)

    print("Loading lyric dataset for translation tasks...")
    lyric_dataset_dict = get_lyric_datasets("../data_pack/multi_lyric.json", train_test_ratio=0.8)
    train_dataset_lyrics = lyric_dataset_dict['train']
    test_dataset_lyrics = lyric_dataset_dict['test']

    # 3. 合并数据(这里为了示例，将math + lyrics合到一起训练)
    #    也可以分阶段先训math再训lyrics
    full_train_data = concatenate_datasets([gsm8k_dataset, train_dataset_lyrics])

    # 4. PPO 配置
    ppo_config = PPOConfig(
        batch_size=16,
        forward_batch_size=4,
        learning_rate=5e-6,
        log_with=None,  # 若你想用wandb，可以设置 log_with="wandb"
        optimize_cuda_graphs=False,
        early_stopping=False,
        num_epochs=1,
        remove_unused_columns=False,
    )

    # 5. 构建 PPOTrainer
    #    这里把我们的 dynamic_reward_func 直接作为一个「自定义Reward」去应用
    ppo_trainer = PPOTrainer(
        config=ppo_config,
        model=model,
        tokenizer=tokenizer
    )

    # 6. 环境 + Agent (如果你希望自己管理数据迭代/生成/奖励计算)
    #    也可以直接使用 ppo_trainer自带的 generate & step等方式，这里更显式地做演示
    environment = TextRLTaskEnvironment(
        dataset=full_train_data,
        tokenizer=tokenizer,
        reward_func=dynamic_reward_func
    )
    agent = TextRLLanguageAgent(model=model, tokenizer=tokenizer)

    # 7. 训练循环 (简化演示)
    num_train_episodes = 100  # 训练回合数
    dataset_size = len(full_train_data)

    # PPOTrainer 自带一个 rollout buffer；这里手动管理一些收集过程
    for episode in range(num_train_episodes):
        # 这里随机抽 batch_size 条数据
        # ppo_config.batch_size 是指 PPO 的一次训练batch大小
        indices = random.sample(range(dataset_size), ppo_config.batch_size)
        prompts, references, meta = environment.reset(indices)

        # 智能体生成动作(完整回复)
        completions = agent.act(prompts)

        # 计算奖励
        rewards, done, info = environment.step(prompts, completions, references, meta)

        # 现在我们要将 prompts + completions(动作) + rewards 交给 ppo_trainer
        # ppo_trainer会做 Advantage 估计、价值函数更新、策略梯度更新
        texts = []
        for i in range(len(prompts)):
            # 把 prompt + completion 拼起来
            prompt_text = ""
            for pm in prompts[i]:
                prompt_text += pm['content'] + "\n"
            completion_text = completions[i][0]['content']
            texts.append(prompt_text + completion_text)

        # PPOTrainer需要 (queries, responses, rewards)
        # queries: prompt部分
        # responses: 完整生成(不包含prompt)
        # 但也可以把 prompts + completions 全部作为 response, 只要保持一致性
        # 这里我们简化： queries = prompt_text, responses = completion_text
        query_texts = [p[-1]['content'] for p in prompts]  # 只拿user问的那部分
        response_texts = [c[0]['content'] for c in completions]

        # ppo step
        train_stats = ppo_trainer.step(
            queries=query_texts,
            responses=response_texts,
            rewards=rewards
        )

        # 你也可以查看 train_stats 内部的一些损失和优势信息
        if (episode+1) % 10 == 0:
            print(f"Episode {episode+1}/{num_train_episodes} done. avg reward={sum(rewards)/len(rewards):.4f}")

    print("=== PPO Training Finished ===")

    # 训练完可以再测试:
    evaluate_model(ppo_trainer, test_dataset_lyrics, tokenizer)


if __name__ == "__main__":
    main()

代码要点说明

1. 环境 (TextRLTaskEnvironment)

   • reset(indices)：根据给定索引，从数据集中取对应样本，返回 prompts（即多轮对话的上下文，如system+user）、references（参考答案/期望输出，用于部分奖励函数），以及 meta（包含额外的信息如task类型）。
   • step(...)：将智能体生成的回复 completions 同 prompts 一起送入 reward_func，得到奖励列表 rewards，并返回 (rewards, done, info)。在这个示例里，一次生成就结束，所以 done 全为 True。

2. 智能体 (TextRLLanguageAgent)

   • 以一个带有策略头(和价值头)的语言模型为底层，接收 prompts 并生成 completions。
   • 这里使用一次性生成文本的方式示例，如果要逐token交互，需要再细化到每一步token的动作和状态，属于更复杂的实现。

3. PPOTrainer

   • 来自 trl 库的核心类。内部会构建一个 rollout buffer 并执行 PPO 算法的更新：
     • 先对当前策略生成的轨迹进行 logprob / value 的前向计算
     • 结合外部给的 rewards 估计优势(Advantage)
     • 利用 PPO 的损失函数 (包括策略损失、价值函数损失、熵正则化等) 来更新模型权重。

4. 优势函数(Advantage) 计算

   • 在 trl 库中自动完成。它会基于每个生成的 step (通常是token-level) 计算 (Q - V) 或 GAE(广义优势估计) 等。
   • 如果你想自己手写，也可以参照 PPO 原理，对每条轨迹的 reward、value 进行： $$A_t = \delta_t + (\gamma\lambda)\delta_{t+1} + ... + ...$$ 其中 $\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$。

5. 动作 (Action)

   • 这里示例中把「生成的回答文本」当做一次整体动作，对于更高粒度(每个token)的操作，可以依赖 trl 提供的token-level PPO流程，但示例中为简化只展示概念。

6. 奖励 (Reward)

   • 由我们自定义的 dynamic_reward_func 返回的标量列表 rewards。
   • 同时也可以将「人工反馈」(Human Feedback) 与「自动打分」相结合，如：

     python
     final_rewards[i] = auto_score + human_score

   • 或者在对话式系统中动态收集用户打分并融入到 rewards。

7. 完整性

   • 上述示例里已经包含了从头到尾的所有代码，包括导入库、定义数据处理函数、奖励函数、环境、智能体以及PPO训练循环，不做任何省略。你可以在此基础上做裁剪或扩展，或切换成 GRPO (只需把 PPOTrainer/PPOConfig 替换成 GRPOTrainer/GRPOConfig 并修改相关超参即可)。

这样，你就获得了一个「更合理的强化学习互动策略」示例：通过环境和智能体的明确定义来分工、利用动作(生成文本)与奖励函数(组合多种评价维度)得到回报、最终通过优势函数(PPO内部计算)来更新策略网络。希望对你有所帮助!