RLHF 101: 深入了解 RLHF 技术

本博客深入探讨了基于人类反馈的强化学习（RLHF）框架的完整训练流程。我们将探索每一个阶段——从数据生成和奖励模型推理到大型语言模型（LLM）的最终训练。我们的目标是通过提供所有必要的代码和所用环境的精确规格，确保一切都完全可复现。在本文结束时，您应该能够使用您选择的算法，通过任何指令数据集来训练任何模型。

前言：设置与环境

本教程将使用以下设置：

• 数据集：UltraFeedback，一个精心策划的通用聊天提示数据集。
• 基础模型：Llama-3-8B-it，最先进的指令微调模型。
• 奖励模型：Armo，一个为评估生成输出而优化的鲁棒奖励模型。
• 训练算法：REBEL，一种专为高效 RLHF 优化而设计的先进算法。

首先，克隆我们的仓库，其中包含本教程所需的所有资源：

git clone https://github.com/ZhaolinGao/REBEL
cd REBEL

我们为管道的不同阶段使用两个独立的环境：

• vllm：处理数据生成，利用高效的 vllm 库。
• rebel：用于训练 RLHF 模型。

您可以使用提供的 YAML 文件安装这两个环境：

conda env create -f ./envs/rebel_env.yml
conda env create -f ./envs/vllm_env.yml

第一部分：数据生成

在本节中，我们将使用 vllm 加载基础模型进行快速推理，准备数据集，并为数据集中的每个提示生成多个响应。本部分的完整代码可在此处获取：https://github.com/ZhaolinGao/REBEL/blob/main/src/ultrafeedback_largebatch/generate.py。

激活 vllm 环境：

conda activate vllm

首先，使用 vllm 加载基础模型和分词器：

from transformers import AutoTokenizer
from vllm import LLM
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
llm = LLM(
    model="meta-llama/Meta-Llama-3-8B-Instruct",
    tensor_parallel_size=8,
)

这里，tensor_parallel_size 指定了要使用的 GPU 数量。

接下来，加载 UltraFeedback 数据集：

from datasets import load_dataset
dataset = load_dataset("allenai/ultrafeedback_binarized_cleaned_train", split='train')

您可以使用 dataset.select 选择数据集的子集。例如，选择前 10,000 行：

dataset = dataset.select(range(10000))

或者，您可以使用 dataset.shard 将数据集分成块，用于像 SPPO 这样的实现，其中每次迭代只在一个块上进行训练。

现在，让我们准备用于生成的数据集。Llama 模型使用特殊标记来区分提示和响应。例如：

<|begin_of_text|><|start_header_id|>user<|end_header_id|>

What is France's capital?<|eot_id|><|start_header_id|>assistant<|end_header_id|>

因此，对于数据集中的每个提示，我们需要在生成之前将其从纯文本转换为这种格式：

def get_message(instruction):
    message = [
        {"role": "user", "content": instruction},
    ]
    return message
prompts = [tokenizer.apply_chat_template(get_message(row['prompt']), tokenize=False, add_generation_prompt=True) for row in dataset]

• get_message 将纯文本提示转换为一个字典，表示它来自用户。
• tokenizer.apply_chat_template 添加所需的特殊标记，并使用 add_generation_prompt=True 在末尾添加响应标记（<|start_header_id|>assistant<|end_header_id|>\n\n）。

最后，我们可以使用 vllm 和刚刚格式化好的提示生成响应。我们将为每个提示生成 5 个响应：

import torch
import random
import numpy as np
from vllm import SamplingParams

def set_seed(seed=5775709):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

for p in range(5):
    set_seed(p * 50)
    sampling_params = SamplingParams(
        temperature=0.8,
        top_p=0.9,
        max_tokens=2048,
        seed=p * 50,
    )
    response = llm.generate(prompts, sampling_params)
    output = list(map(lambda x: x.outputs[0].text, response))
    dataset = dataset.add_column(f"response_{p}", output)

• temperature=0.8, top_p=0.9 是控制生成多样性的常用设置。
• set_seed 用于确保可复现性，并为每个响应设置不同的种子。
• llm.generate 生成响应，结果通过 dataset.add_column 添加到数据集中。

您可以使用以下命令运行完整脚本：

python ./src/ultrafeedback_largebatch/generate.py --world_size NUM_GPU --output_repo OUTPUT_REPO

第二部分：奖励模型推理

在本部分中，我们将计算第一部分中生成的响应的奖励分数。本部分的完整代码可在 此处 获取。

激活 rebel 环境：

conda activate rebel

首先，我们将初始化 Armo 奖励模型管道。此奖励模型是一个经过微调的序列分类模型，根据对话质量为给定对话分配一个标量奖励分数。实现如下：

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from typing import Dict, List

class ArmoRMPipeline:
    def __init__(self, model_id, device_map="cuda", torch_dtype=torch.bfloat16, truncation=True, trust_remote_code=False, max_length=4096):
        self.model = AutoModelForSequenceClassification.from_pretrained(
            model_id,
            device_map=device_map,
            trust_remote_code=trust_remote_code,
            torch_dtype=torch_dtype,
        )
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_id,
            use_fast=True,
        )
        self.truncation = truncation
        self.device = self.model.device
        self.max_length = max_length

    def __call__(self, messages: List[Dict[str, str]]) -> Dict[str, float]:
        input_ids = self.tokenizer.apply_chat_template(
            messages,
            return_tensors="pt",
            padding=True,
            truncation=self.truncation,
            max_length=self.max_length,
        ).to(self.device)
        with torch.no_grad():
            output = self.model(input_ids)
            score = output.score.float().item()
        return score

rm = ArmoRMPipeline("RLHFlow/ArmoRM-Llama3-8B-v0.1", trust_remote_code=True)

现在，我们可以收集奖励分数：

def get_message(instruction, response):
    return [{"role": "user", "content": instruction}, {"role": "assistant", "content": response}]

rewards = {}
for i in range(5):
    rewards[f"response_{i}_reward"] = []
    for row in dataset:
        reward = rm(get_message(row['prompt'], row[f'response_{i}']))
        rewards[f"response_{i}_reward"].append(reward)
for k, v in rewards.items():
    dataset = dataset.add_column(k, v)

• get_message 将用户提示和助手响应格式化为字典列表。
• rm 计算数据集中每个响应的奖励分数。

您可以使用以下命令运行完整脚本：

python ./src/ultrafeedback_largebatch/rank.py --input_repo INPUT_REPO

• INPUT_REPO 是从第一部分保存的仓库，其中包含生成的响应。

第三部分：过滤与标记化

在本部分中，我们将介绍准备训练数据集的过程，包括过滤过长的提示和响应以防止内存不足（OOM），选择最佳和最差响应进行训练，以及删除重复响应。本部分的完整代码可在此处获取：https://github.com/ZhaolinGao/REBEL/blob/main/src/ultrafeedback_largebatch/filter_tokenize.py。

我们首先初始化两个不同的分词器，一个从右侧填充，一个从左侧填充：

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
tokenizer.add_special_tokens({"pad_token": "[PAD]"})
tokenizer_left = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", padding_side='left')
tokenizer_left.add_special_tokens({"pad_token": "[PAD]"})

由于 Llama 没有填充标记，我们手动使用 tokenizer.add_special_tokens 添加它。

这两个不同的分词器允许我们将提示从左侧填充，响应从右侧填充，使它们在中间相遇。通过将左填充提示与右填充响应结合，我们确保：

• 提示和响应在一致位置相遇。
• 相对位置嵌入对于模型训练保持正确。

这是一个示例格式：

[PAD] ... [PAD] <|begin_of_text|><|start_header_id|>user<|end_header_id|>

PROMPT<|eot_id|><|start_header_id|>assistant<|end_header_id|>


RESPONSE<|eot_id|>[PAD] ... [PAD]

我们希望确保

[PAD] ... [PAD] <|begin_of_text|><|start_header_id|>user<|end_header_id|>

PROMPT<|eot_id|><|start_header_id|>assistant<|end_header_id|>

对于所有提示，长度相同，并且

RESPONSE<|eot_id|>[PAD] ... [PAD]

对于所有响应，长度相同。

接下来，我们过滤掉长度超过 1,024 个标记的提示和长度超过 2,048 个标记的响应。过滤过程使用一个辅助函数来创建消息模板：

def get_message(instruction=None, response=None):

    assert instruction != None or response != None

    if response == None:
        message = [
            {"role": "user", "content": instruction},
        ]
    elif instruction == None:
        message = [
            {"role": "assistant", "content": response}
        ]
    else:
        message = [
            {"role": "user", "content": instruction},
            {"role": "assistant", "content": response}
        ]

    return message

dataset = dataset.filter(lambda row: tokenizer.apply_chat_template(get_message(row['prompt']), tokenize=True, add_generation_prompt=True, return_tensors='pt').shape[-1] <= 1024)
for i in range(5):
    dataset = dataset.filter(lambda row: tokenizer.apply_chat_template(get_message(response=row[f'response_{i}']), tokenize=True, add_generation_prompt=False, return_tensors='pt')[:, 5:].shape[-1] <= 2048)

请注意，我们在计算长度时跳过响应的前五个标记，以排除特殊标记（例如 <|begin_of_text|><|start_header_id|>assistant<|end_header_id|>\n\n），并且只计算响应的实际长度加上末尾的 EOS 标记（<|eot_id|>）。

现在我们可以对提示进行左填充标记化，最大长度为 1,024 个标记：

llama_prompt_tokens = []
for row in dataset:
    llama_prompt_token = tokenizer_left.apply_chat_template(
            get_message(row['prompt']), 
            add_generation_prompt=True,
            tokenize=True,
            padding='max_length',
            max_length=1024,
    )
    assert len(llama_prompt_token) == 1024
    assert (llama_prompt_token[0] == 128000 or llama_prompt_token[0] == 128256) and llama_prompt_token[-1] == 271
    llama_prompt_tokens.append(llama_prompt_token)
dataset = dataset.add_column("llama_prompt_tokens", llama_prompt_tokens)

断言用于确保长度始终为 1,024，并且标记化的提示以 [pad] 标记或 <|begin_of_text|> 标记开头，并以 \n\n 标记结尾。

然后，我们选择每个提示奖励最高和最低的响应作为选择和拒绝响应，并对其进行右填充标记化：

chosen, reject, llama_chosen_tokens, llama_reject_tokens, chosen_reward, reject_reward = [], [], [], [], [], []

for row in dataset:

    all_rewards = [row[f"response_{i}_reward"] for i in range(5)]
    chosen_idx, reject_idx = np.argmax(all_rewards), np.argmin(all_rewards)

    chosen.append(row[f"response_{chosen_idx}"])
    reject.append(row[f"response_{reject_idx}"])

    llama_chosen_token = tokenizer.apply_chat_template(
            get_message(response=row[f"response_{chosen_idx}"]),
            add_generation_prompt=False,
            tokenize=True,
            padding='max_length',
            max_length=2048+5,
    )[5:]
    llama_chosen_tokens.append(llama_chosen_token)
    chosen_reward.append(row[f"response_{chosen_idx}_reward"])
    assert len(llama_chosen_token) == 2048
    assert llama_chosen_token[-1] == 128009 or llama_chosen_token[-1] == 128256

    llama_reject_token = tokenizer.apply_chat_template(
            get_message(response=row[f"response_{reject_idx}"]),
            add_generation_prompt=False,
            tokenize=True,
            padding='max_length',
            max_length=2048+5,
    )[5:]
    llama_reject_tokens.append(llama_reject_token)
    reject_reward.append(row[f"response_{reject_idx}_reward"])
    assert len(llama_reject_token) == 2048
    assert llama_reject_token[-1] == 128009 or llama_reject_token[-1] == 128256

dataset = dataset.add_column("chosen", chosen)
dataset = dataset.add_column("chosen_reward", chosen_reward)
dataset = dataset.add_column("llama_chosen_tokens", llama_chosen_tokens)
dataset = dataset.add_column("reject", reject)
dataset = dataset.add_column("reject_reward", reject_reward)
dataset = dataset.add_column("llama_reject_tokens", llama_reject_tokens)

同样，断言用于确保标记化响应的长度始终为 2,048，并且标记化响应以 [pad] 标记或 <|eot_id|> 标记结尾。

最后，我们过滤掉选择和拒绝响应相同的数据行：

dataset = dataset.filter(lambda row: row['chosen'] != row['reject'])

并将数据集分为训练集和包含 1,000 个提示的测试集：

dataset = dataset.train_test_split(test_size=1000, shuffle=True)

您可以使用以下命令运行完整脚本：

python ./src/ultrafeedback_largebatch/filter_tokenize.py --input_repo INPUT_REPO

• INPUT_REPO 是从第二部分保存的仓库，其中包含每个响应的奖励。

第四部分：使用 REBEL 进行训练

在 REBEL 的每次迭代 $tt$ 中，我们旨在解决以下平方损失回归问题： 其中 $\eta \backslash eta$ 是一个超参数， $\theta \backslash theta$ 是模型的参数， $xx$ 是提示， ${\mathcal{D}}_t\backslash mathcal\{D\}\_t$ 是我们从前三部分收集到的数据集， $yy$ 和 $y^{\mathrm{\prime }}y\text{'}$ 是 $xx$ 的响应， ${\pi }_{\theta }(y\mathrm{\mid }x)\backslash pi\_\backslash theta(y|x)$ 是在参数化策略 ${\pi }_{\theta }\backslash pi\_\backslash theta$ 下给定提示 $xx$ 生成响应 $yy$ 的概率，而 $r(x,y)r(x,\; y)$ 是响应 $yy$ 对于提示 $xx$ 的奖励，该奖励从第二部分获得。

在本教程中，我们使用基础模型 ${\pi }_{{\theta }_0}\backslash pi\_\{\backslash theta\_0\}$ 演示 REBEL 的单次迭代（t=0）。对于多迭代训练，您可以重复第一部分到第四部分，每次迭代都使用前一次迭代中训练好的模型进行初始化。

本部分的完整代码可在 此处 获取。为了使用 8 个 GPU 进行全参数训练，我们使用 Accelerate 库和 Deepspeed Stage 3，运行命令如下：

accelerate launch --config_file accelerate_cfgs/deepspeed_config_stage_3.yaml --main-process-port 29080 --num_processes 8 src/ultrafeedback_largebatch/rebel.py --task.input_repo INPUT_REPO --output_dir OUTPUT_DIR

• INPUT_REPO 是从第三部分保存的仓库，其中包含标记化的提示和响应。
• OUTPUT_DIR 是保存模型的目录。

步骤 1：初始化和加载

我们首先初始化分布式训练的批次大小：

args.world_size = accelerator.num_processes
args.batch_size = args.world_size * args.per_device_train_batch_size * args.gradient_accumulation_steps
args.local_batch_size = args.per_device_train_batch_size * args.gradient_accumulation_steps
args.rebel.num_updates = args.total_episodes // args.batch_size

• args.world_size 是我们正在使用的 GPU 数量。
• args.local_batch_size 是每个 GPU 的批次大小。
• args.batch_size 是实际的训练批次大小。
• args.rebel.num_updates 是要执行的总更新次数，args.total_episodes 是要训练的数据点数量。通常，我们将 args.total_episodes 设置为训练集在一个 epoch 中的大小。

接下来，我们加载模型和分词器，确保禁用 dropout 层：

tokenizer = AutoTokenizer.from_pretrained(
                args.base_model, 
                padding_side='right',
                trust_remote_code=True,
            )
tokenizer.add_special_tokens({"pad_token": "[PAD]"})
policy = AutoModelForCausalLM.from_pretrained(
            args.base_model,
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            attn_implementation="flash_attention_2",
        )
disable_dropout_in_model(policy)

步骤 2：记录概率

为了方便超参数调优（例如，学习率和 $\eta \backslash eta$），我们预先计算并保存 ${\pi }_{{\theta }_t}\backslash pi\_\{\backslash theta\_t\}$ 的对数概率。因此，当我们第一次运行以下代码时，try 将失败，因为数据集中尚未提供 chosen_logprob 或 reject_logprob。

compute_log = False
try:
    dataset = load_dataset(args.task.input_repo + '_logprob', split='train')
    dataset = dataset.with_format("torch", columns=["llama_prompt_tokens", 
                                                    "llama_chosen_tokens", "chosen_reward", "chosen_logprob",
                                                    "llama_reject_tokens", "reject_reward", "reject_logprob"])
    temp_dataloader = DataLoader(dataset, batch_size=args.local_batch_size, shuffle=True)
    validation_dataset = load_dataset(args.task.input_repo + '_logprob', split='test')
    validation_dataset = validation_dataset.with_format("torch", columns=["llama_prompt_tokens", 
                                                        "llama_chosen_tokens", "chosen_reward", "chosen_logprob",
                                                        "llama_reject_tokens", "reject_reward", "reject_logprob"])
except:
    dataset = load_dataset(args.task.input_repo, split='train')
    dataset = dataset.with_format("torch", columns=["llama_prompt_tokens", 
                                                    "llama_chosen_tokens", "chosen_reward",
                                                    "llama_reject_tokens", "reject_reward"])
    temp_dataloader = DataLoader(dataset, batch_size=args.local_batch_size, shuffle=True)
    validation_dataset = load_dataset(args.task.input_repo, split='test')
    validation_dataset = validation_dataset.with_format("torch", columns=["llama_prompt_tokens", 
                                                        "llama_chosen_tokens", "chosen_reward",
                                                        "llama_reject_tokens", "reject_reward"])
    compute_log = True

在这里，我们需要一个临时数据加载器 temp_dataloader，在计算对数概率之前使用 accelerator.prepare 准备模型，以实现多 GPU 推理。

现在，我们可以计算并保存对数概率：

def gather_logprob(args, model, tokenizer, query, response, device):

    query_response = torch.cat((query, response), dim=-1).long().to(device).unsqueeze(0)
    response = response.long().to(device).unsqueeze(0)
    attention_mask = query_response != tokenizer.pad_token_id
    input_ids = torch.masked_fill(query_response, ~attention_mask, tokenizer.eos_token_id)
    with torch.no_grad():
        output = model(
                    input_ids=input_ids, 
                    attention_mask=attention_mask,
                    return_dict=True,
                 )
        logits = output.logits[:, args.task.maxlen_prompt - 1 : -1]
        logits /= args.task.temperature + 1e-7
        all_logprob = F.log_softmax(logits, dim=-1)
        logprob = torch.gather(all_logprob, 2, input_ids[:, args.task.maxlen_prompt:].unsqueeze(-1)).squeeze(-1)
        sequence_length = first_true_indices(response == tokenizer.pad_token_id) - 1
        seq_mask = torch.arange(args.task.maxlen, device=device).unsqueeze(0).expand_as(response) <= sequence_length.unsqueeze(1)
        
        return (logprob * seq_mask).sum(-1)


def gather_all_logprob(args, process_idx, policy, tokenizer, dataset, device):

    batch_size = len(dataset) // args.world_size + 1
    start_idx = batch_size * process_idx

    # make batch size same for accelerator.gather
    if start_idx + batch_size > len(dataset):
        start_idx = len(dataset) - batch_size

    chosen_logprob, reject_logprob, index = [], [], []

    with torch.no_grad():
        for i in tqdm(range(start_idx, start_idx + batch_size)):

            chosen_logprob.append(gather_logprob(args, policy, tokenizer, dataset[i]["llama_prompt_tokens"], dataset[i]["llama_chosen_tokens"], device))
            reject_logprob.append(gather_logprob(args, policy, tokenizer, dataset[i]["llama_prompt_tokens"], dataset[i]["llama_reject_tokens"], device))
            index.append(i)

        chosen_logprob = torch.cat(chosen_logprob)
        reject_logprob = torch.cat(reject_logprob)
        index = torch.LongTensor(index).to(device)

    chosen_logprob = accelerator.gather(chosen_logprob).cpu().tolist()
    reject_logprob = accelerator.gather(reject_logprob).cpu().tolist()
    index = accelerator.gather(index).cpu().tolist()

    chosen_logprobs = [0] * len(dataset)
    reject_logprobs = [0] * len(dataset)

    for i, data_i in enumerate(index):
        chosen_logprobs[data_i] = chosen_logprob[i]
        reject_logprobs[data_i] = reject_logprob[i]
        
    return chosen_logprobs, reject_logprobs

if compute_log:
    accelerator.print('gathering validation logprob')
    chosen_logprob, reject_logprob = gather_all_logprob(args, accelerator.process_index, accelerator.unwrap_model(policy), tokenizer, validation_dataset, device)
    validation_dataset = validation_dataset.add_column("chosen_logprob", chosen_logprob)
    validation_dataset = validation_dataset.add_column("reject_logprob", reject_logprob)
    validation_dataset = validation_dataset.with_format("torch", columns=["llama_prompt_tokens", 
                                                                            "llama_chosen_tokens", "chosen_reward", "chosen_logprob",
                                                                            "llama_reject_tokens", "reject_reward", "reject_logprob"])

    accelerator.print('gathering logprob')
    chosen_logprob, reject_logprob = gather_all_logprob(args, accelerator.process_index, accelerator.unwrap_model(policy), tokenizer, dataset, device)
    dataset = dataset.add_column("chosen_logprob", chosen_logprob)
    dataset = dataset.add_column("reject_logprob", reject_logprob)
    dataset = dataset.with_format("torch", columns=["llama_prompt_tokens", 
                                                    "llama_chosen_tokens", "chosen_reward", "chosen_logprob",
                                                    "llama_reject_tokens", "reject_reward", "reject_logprob"])
    if accelerator.is_main_process:
        temp = DatasetDict({
            "train" : dataset,
            "test"  : validation_dataset,
        })
        temp.push_to_hub(args.task.input_repo + '_logprob')

• gather_logprob 计算数据集中每个提示和响应对的对数概率。
• gather_all_logprob 通过将数据集划分为 args.world_size 个大小相同的段来并行计算。如果数据集大小不能被段数完美划分，我们让最后一个段与前一个段重叠。大小相同的段确保 accelerator.gather 能够正常工作。

步骤 3：训练

再看看 REBEL 目标，我们现在唯一需要训练的就是计算 ${\pi }_{\theta }(y\mathrm{\mid }x)\backslash pi\_\backslash theta(y|x)$。我们可以通过以下方式计算：

output = policy(
    input_ids=input_ids, 
    attention_mask=attention_mask,
    return_dict=True,
    output_hidden_states=True,
)
logits = output.logits[:, args.task.maxlen_prompt - 1 : -1]
logits /= args.task.temperature + 1e-7
new_all_logprobs = F.log_softmax(logits, dim=-1)
new_logprobs = torch.gather(new_all_logprobs, 2, input_ids[:, args.task.maxlen_prompt:].unsqueeze(-1)).squeeze(-1)
new_logprobs = (new_logprobs * mb_seq_mask).sum(-1)

• output.logits 包含 input_ids 序列中所有词汇标记的 logits。
• output.logits[:, args.task.maxlen_prompt - 1 : -1] 是仅用于响应序列中所有词汇标记的 logits。它偏移了 1，因为位置 $pp$ 处的 logits 指的是位置 $p+1p+1$ 处的 logits。
• 我们将 logits 除以 args.task.temperature，以获得生成期间的实际概率。
• torch.gather 用于收集响应中的对应标记。
• mb_seq_mask 掩盖了填充。

步骤 4：损失计算

最后，我们可以计算损失：

ratio_logprob = new_logprobs - mb_logprobs
ratio_logprob = ratio_logprob[:args.per_device_train_batch_size] - ratio_logprob[args.per_device_train_batch_size:]
reg_diff = ratio_logprob - args.rebel.eta * (mb_chosen_reward - mb_reject_reward)
loss = (reg_diff ** 2).mean()

性能

仅通过上述 4 个部分的迭代一次，我们就可以大幅提升基础模型在 AlpacaEval、MT-Bench 和 ArenaHard 上的性能：

总结

在本文中，我们概述了 RLHF 的管道，涵盖了从数据生成到实际训练的整个过程。虽然我们特别关注 REBEL 算法，但此管道用途广泛，可以轻松适应其他方法，例如 DPO 或 SimPO。除具体的损失公式外，这些方法所需的组件已包含在内。

如果您觉得此实现有用，请考虑引用我们的工作：

@misc{gao2024rebel,
      title={REBEL: Reinforcement Learning via Regressing Relative Rewards}, 
      author={Zhaolin Gao and Jonathan D. Chang and Wenhao Zhan and Owen Oertell and Gokul Swamy and Kianté Brantley and Thorsten Joachims and J. Andrew Bagnell and Jason D. Lee and Wen Sun},
      year={2024},
      eprint={2404.16767},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}