RolloutWorkflow 参考

本文档描述 RolloutWorkflow 抽象，它是 AReaL 强化学习流水线中实现 Rollout 生成的核心接口。

注意：

1. 本页面向希望深入理解代码库的开发者。对于代理 RL 训练，请使用 Agentic RL Guide 中描述的高级 API。

2. 遗留模式：直接子类化 RolloutWorkflow 被视为遗留方式，不应主动使用。对于新的代理 RL 工作流，请使用带有 async def run() 的代理工作流模式。

概述

RolloutWorkflow 定义如何从输入数据生成训练轨迹。它封装了以下逻辑：

• 对提示进行分词并准备模型输入

• 调用推理引擎生成回复

• 为生成的输出计算奖励

• 将结果打包成用于训练的张量字典

接口

from areal.api.workflow_api import RolloutWorkflow

class RolloutWorkflow(ABC):
    @abstractmethod
    async def arun_episode(
        self, engine: InferenceEngine, data: dict[str, Any]
    ) -> dict[str, Any] | None | dict[str, InteractionWithTokenLogpReward]:
        """运行单个工作流 episode。"""
        ...

参数

参数	类型	描述
engine	InferenceEngine	用于生成模型回复的推理引擎
data	dict[str, Any]	数据加载器的单个样本

返回类型

arun_episode 方法支持三种返回类型：

返回类型	描述
dict[str, torch.Tensor]	用于训练的标准张量格式
dict[str, InteractionWithTokenLogpReward]	Token 级别的交互（自动转换为张量）；由高级 ArealOpenAI API 生成
None	拒绝的轨迹，排除在训练之外

张量字典格式

返回张量字典时，预期包含以下字段：

字段	形状	类型	必需	描述
input_ids	[batch_size, seq_len]	int32	是	Token ID（提示 + 生成内容）
attention_mask	[batch_size, seq_len]	bool	是	有效 token 掩码
loss_mask	[batch_size, seq_len]	int32	否	生成内容 token 掩码（1 = 训练）
logprobs	[batch_size, seq_len]	float32	否	每个 token 的对数概率
rewards	[batch_size]	float32	否	每序列奖励
versions	[batch_size, seq_len]	int32	否	生成 token 时的权重版本

返回值示例：

return {
    "input_ids": torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.int32),
    "attention_mask": torch.ones(1, 5, dtype=torch.bool),
    "loss_mask": torch.tensor([[0, 0, 1, 1, 1]], dtype=torch.int32),
    "logprobs": torch.tensor([[0.0, 0.0, -0.5, -0.3, -0.2]], dtype=torch.float32),
    "rewards": torch.tensor([1.0], dtype=torch.float32),
    "versions": torch.tensor([[0, 0, 1, 1, 1]], dtype=torch.int32),
}

工作流上下文

在 arun_episode 中，通过 workflow_context 模块访问执行上下文。每个工作流实例有其自己的隔离上下文：

from areal.infra import workflow_context

async def arun_episode(self, engine, data):
    # 获取当前执行上下文
    ctx = workflow_context.get()

    # 检查是否在评估模式运行
    if ctx.is_eval:
        # 使用不同参数进行评估
        ...

    # 获取用于日志的任务 ID
    task_id = ctx.task_id

    # 根据模式获取统计作用域（"rollout" 或 "eval-rollout"）
    scope = workflow_context.stat_scope()

轨迹转储

当 InferenceEngineConfig.dump_to_file=True 时，轨迹自动保存到磁盘用于调试和分析。

配置

rollout:
  dump_to_file: true
  fileroot: "/path/to/logs"
  tokenizer_path: "model/tokenizer" # 文本解码必需

输出位置

轨迹保存到：

{fileroot}/{experiment_name}/{trial_name}/[rollout|eval-rollout]/{version}/{task_id}.jsonl

示例：

/tmp/areal/my_exp/trial1/rollout/5/42.jsonl

输出格式

JSONL 文件的每一行包含：

{
  "task_id": 42,
  "sample_idx": 0,
  "seqlen": 512,
  "prompt_len": 128,
  "head_version": 5,
  "tail_version": 6,
  "version_rle": [[5, 100], [6, 200]],
  "reward": 1.0,
  "prompt": "<|im_start|>user\nWhat is 2+2?<|im_end|>\n<|im_start|>assistant\n",
  "completion": "The answer is 4.<|im_end|>"
}

字段说明：

字段	说明
task_id	Batch 任务标识符
sample_idx	Batch 内的样本索引
seqlen	有效序列长度
prompt_len	第一个生成 token 的位置（mask.index(1)）。多轮 agent rollout 中为首次 assistant 生成位置，而非 seqlen - sum(mask)
head_version	Per-sample 最小模型版本，仅统计 loss_mask==1 的 token
tail_version	Per-sample 最大模型版本，仅统计 loss_mask==1 的 token
version_rle	RLE 编码的 per-token 版本序列（仅输出 token），例如 [[5, 100], [6, 200]]
reward	该样本的 reward 值
prompt	解码后的 prompt 文本
completion	解码后的 completion 文本
segments	（仅多轮） 分段列表 [{"role": "prompt"|"gen"|"context", "len": N, "text": "..."}]

目录命名： {version} 目录以 batch 内的全局最大版本（global_tail）命名。 同一目录内的单条 record 的 tail_version 可能 <= 目录名。

分组 Rollout

分组 Rollout 对每个输入提示多次运行相同工作流，生成多样化回复用于训练。这对于像 GRPO 这样受益于每个提示多个样本的算法很有用。

配置

提交 Rollout 时设置 group_size：

engine.submit(
    data=sample,
    workflow=MyWorkflow,
    workflow_kwargs={...},
    group_size=4,  # 每个输入运行工作流 4 次
)

或通过 CLI：

rollout:
  group_size: 4

工作原理

当 group_size > 1 时，工作流被包装在 GroupedRolloutWorkflow 中：

1. 包装器使用 asyncio.gather 并发运行 arun_episode group_size 次

2. 根据类型合并结果：

   • 张量字典：沿批次维度连接

   • InteractionWithTokenLogpReward 字典：合并为单个字典

3. 每个 slot 可用时返回正常结果类型，不可用时返回 None。None 有意保持不透明；原因分类与重试策略仍由 producer 负责。

4. 包装器只等待最初提交的 slots，既不会重试不可用 slot，也不会复制可用结果。

5. 可用 slots 会各保留一次并连接；实际数量会在 reward 和 advantage normalization 中继续作为 prompt-group 边界。

6. 设置 reward_normalization=True 时，group 通过最小大小检查后，只对可用 rollout 的 interaction rewards 进行归一化。drop_incomplete_group=True 仍要求所有原始 slot 成功；若保留的 rollout 缺少某一行的 reward，整个 group 会被丢弃。

7. min_usable_group_size 默认为 1。当 RL trainer 的 reward 或 advantage normalization 使用 group statistics（该统计量至少需要两个观测值）时，才会将其设为 2；singleton 目标 group （n_samples: 1）本身即是完整的，因此下限保持为 1。设置 actor.min_usable_group_size 可覆盖该推导值；使用 group statistics 时，低于 2 的显式值会被拒绝。低于下限的 group 返回 None，异步 collector 随后会接收另一个已就绪的 prompt group。采用 batch-relative PPO 或 REINFORCE 时则会保留可用的 singleton。

8. v1 每次 arun_episode 调用对应一个逻辑 rollout。上下文压缩或 agent.export_style: individual 可以导出多行，group 不会因此中止。Collector 在轨迹的 rollout_group 字段中保存 RolloutGroup，记录每个 rollout 的连续行数和可选 reward 参考值。Batch 合并将其移入 TrajBatchMeta，拆分时再恢复。可用 group 大小按逻辑 rollout 计数，n_samples: 1 也遵循该规则。

9. Group 和 batch reward normalization 都对每个逻辑 rollout 使用一个参考值。若同一 rollout 各行 reward 不同，workflow 必须在张量字典中或导出的 InteractionWithTokenLogpReward 上 提供有限标量 rollout_reward；同一 rollout 中所有显式参考值必须一致。省略时，只能从相等的 行 reward 取得参考值。每一行保留自己的 reward，并使用相同的平移和缩放；leave-one-out 基线排除整个逻辑 rollout。当算出的参考值标准差不大于 normalization epsilon 时，保留平移并使用除数 1。内置 v1 agent workflow 为 individual 导出显式提供最终 reward 作为参考值，同时保留各行折扣后的 reward 和已提供的参考值。自定义 workflow 必须为不同行 reward 声明自己的参考值，不会自动猜测最终值、总和或均值。

10. 内置按行长度计算的 overlong penalty 不改变显式参考值；actor 的 reward bias、scaling 和 clipping 同时作用于行 reward 与参考值。未提供显式参考值时，施加长度惩罚后的行 reward 仍须相等。Advantage normalization 保持现有 masked token 统计及按 token 的 leave-one-out 行为，逻辑计数仅用于 singleton 回退。GAE 仍分别在各行上计算。

PPO 系列 actor loss 默认仍按全局 token 加权。因此，有更多有效 response tokens 的 partial group 会比更小或更短的 group 获得更高 loss weight。这是保持向后兼容的现有 estimator，并不表示各 prompt 隐式等权。

Grouped rollout 会导出 target_slot_count、usable_slot_count、 trainable_slot_count、fully_masked_group、singleton_slot_group、 pre_filter_usable_slot_yield 和 pre_filter_trainable_slot_yield。这些指标统计最初的 rollout 调用；should_accept_fn 运行后的最终 accepted 和 rejected 数量仍由 collector metrics 提供。PPO 训练会另外报告物理 usable group size 与有效 token 的 loss-weight 分布，其中包括按 size 区分的 group_loss_weight_size_<N> 指标。

输出形状

当 group_size=4、工作流返回 [1, seq_len] 张量且 4 个 slot 均可用时，分组输出的形状为 [4, seq_len]。如果接受的是不完整 group，则第一维是实际可用的 slot 数量。

实现

来自 areal/infra/remote_inf_engine.py：

class GroupedRolloutWorkflow(RolloutWorkflow):
    async def arun_episode(self, engine, data):
        # 并发运行 N 次
        results = await asyncio.gather(
            *[self.workflow.arun_episode(engine, data)
              for _ in range(self.group_size)]
        )

        # 正常结果表示可用，None 表示不可用
        valid_results = [r for r in results if r is not None]
        if len(valid_results) < self.min_usable_group_size:
            return None

        # 根据结果类型合并
        if all_interaction_dicts(valid_results):
            return merge_dicts(valid_results)
        else:
            return concat_padded_tensors(valid_results)

实现自定义工作流

创建自定义工作流：

1. 子类化 RolloutWorkflow：

from areal.api.workflow_api import RolloutWorkflow

class MyWorkflow(RolloutWorkflow):
    def __init__(self, tokenizer, gconfig, **kwargs):
        self.tokenizer = tokenizer
        self.gconfig = gconfig

    async def arun_episode(self, engine, data):
        # 1. 准备输入
        input_ids = self.tokenizer.encode(data["prompt"])

        # 2. 生成回复
        req = ModelRequest(
            rid=uuid.uuid4().hex,
            input_ids=input_ids,
            gconfig=self.gconfig,
            tokenizer=self.tokenizer,
        )
        resp = await engine.agenerate(req)

        # 3. 计算奖励
        reward = self.compute_reward(resp, data)

        # 4. 返回张量字典（或 None 拒绝）
        if reward < 0:
            return None

        return self.build_tensor_dict(resp, reward)

2. 向训练器注册：

trainer.train(
    workflow=MyWorkflow,
    workflow_kwargs={
        "tokenizer": tokenizer,
        "gconfig": config.gconfig,
    },
)

工作流解析

工作流可以通过多种方式指定：

格式	示例	描述
实例	MyWorkflow(...)	预实例化的工作流
类	MyWorkflow	类（需要 kwargs）
字符串路径	"my_module.MyWorkflow"	动态导入
代理工作流	任何带 async def run() 的类	带代理支持封装

训练系统自动将这些解析为 RolloutWorkflow 实例。

另见

• Agentic RL Tutorial - 使用代理框架训练

• Adding Custom Workflows - 分步指南