RolloutWorkflow Reference

This document describes the RolloutWorkflow abstraction, the core interface for implementing rollout generation in AReaL’s reinforcement learning pipeline.

Notes:

1. This page targets developers seeking a deep understanding of the codebase. For agentic RL training, use the high-level API described in the Agentic RL Guide.

2. Legacy pattern: Directly subclassing RolloutWorkflow is considered legacy and should not be used proactively. For new agentic RL workflows, use the agent workflow pattern with async def run() instead.

Overview

A RolloutWorkflow defines how to generate training trajectories from input data. It encapsulates the logic for:

• Tokenizing prompts and preparing model inputs

• Calling the inference engine to generate completions

• Computing rewards for generated outputs

• Packaging results into tensor dictionaries for training

Interface

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]:
        """Run a single episode of the workflow."""
        ...

Parameters

Parameter	Type	Description
engine	InferenceEngine	Inference engine for generating model responses
data	dict[str, Any]	A single sample from the dataloader

Return Types

The arun_episode method supports three return types:

Return Type	Description
dict[str, torch.Tensor]	Standard tensor format for training
dict[str, InteractionWithTokenLogpReward]	Token-level interactions (auto-converted to tensors); produced by the high-level ArealOpenAI API
None	Rejected trajectory, excluded from training

Tensor Dictionary Format

When returning a tensor dictionary, the following fields are expected:

Field	Shape	Type	Required	Description
input_ids	[batch_size, seq_len]	int32	Yes	Token IDs (prompt + completion)
attention_mask	[batch_size, seq_len]	bool	Yes	Valid token mask
loss_mask	[batch_size, seq_len]	int32	No	Completion token mask (1 = train)
logprobs	[batch_size, seq_len]	float32	No	Log probabilities per token
rewards	[batch_size]	float32	No	Per-sequence rewards
versions	[batch_size, seq_len]	int32	No	Weight version when token generated

Example return value:

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),
}

Workflow Context

Inside arun_episode, access the execution context via the workflow_context module. Each workflow instance has its own isolated context:

from areal.infra import workflow_context

async def arun_episode(self, engine, data):
    # Get current execution context
    ctx = workflow_context.get()

    # Check if running in evaluation mode
    if ctx.is_eval:
        # Use different parameters for evaluation
        ...

    # Get task ID for logging
    task_id = ctx.task_id

    # Get stats scope based on mode ("rollout" or "eval-rollout")
    scope = workflow_context.stat_scope()

Trajectory Dumping

When InferenceEngineConfig.dump_to_file=True, trajectories are automatically saved to disk for debugging and analysis.

Configuration

rollout:
  dump_to_file: true
  fileroot: "/path/to/logs"
  tokenizer_path: "model/tokenizer" # Required for text decoding

Output Location

Trajectories are saved to:

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

Example:

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

Output Format

Each line in the JSONL file contains:

{
  "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|>"
}

Field descriptions:

Field	Description
task_id	Batch task identifier
sample_idx	Index of the sample within the batch
seqlen	Effective sequence length
prompt_len	Index of the first generated token (mask.index(1)). For multi-turn agent rollouts this is the position of the first assistant generation, not seqlen - sum(mask)
head_version	Per-sample minimum model version among loss_mask==1 tokens
tail_version	Per-sample maximum model version among loss_mask==1 tokens
version_rle	Run-length encoded per-token version sequence (output tokens only), e.g. [[5, 100], [6, 200]]
reward	Reward value for this sample
prompt	Decoded prompt text
completion	Decoded completion text
segments	(multi-turn only) List of {"role": "prompt"|"gen"|"context", "len": N, "text": "..."}

Directory naming: The {version} directory is named by the batch-global maximum version (global_tail). Individual records within the same directory may have tail_version <= global_tail.

Grouped Rollout

Grouped rollout runs the same workflow multiple times per input prompt, producing diverse completions for training. This is useful for algorithms like GRPO that benefit from multiple samples per prompt.

Configuration

Set group_size when submitting rollouts:

engine.submit(
    data=sample,
    workflow=MyWorkflow,
    workflow_kwargs={...},
    group_size=4,  # Run workflow 4 times per input
)

Or via CLI:

rollout:
  group_size: 4

How It Works

When group_size > 1, the workflow is wrapped in GroupedRolloutWorkflow:

1. The wrapper runs arun_episode concurrently group_size times using asyncio.gather

2. Results are merged based on their type:

   • Tensor dictionaries: Concatenated along the batch dimension

   • InteractionWithTokenLogpReward dicts: Merged into a single dictionary

3. Each slot returns its normal result type when usable and None when unusable. None is intentionally opaque: classification and retry policy stay in the producer.

4. The wrapper waits only for the original slots. It neither retries an unusable slot nor duplicates a usable result.

5. Usable slots are retained exactly once and concatenated. Their actual count remains a prompt-group boundary during reward and advantage normalization.

6. With reward_normalization=True, interaction rewards are normalized across the usable rollouts after the group passes the minimum-size filter. drop_incomplete_group=True still requires every original slot to succeed. A surviving rollout with a missing row reward causes the group to be dropped.

7. min_usable_group_size defaults to 1. The RL trainer sets it to 2 when reward or advantage normalization uses group statistics, because that statistic needs at least two observations; a singleton target group (n_samples: 1) is complete by definition and keeps the minimum of 1. Setting actor.min_usable_group_size replaces this derived value; explicit values below 2 are rejected while group statistics are in use. Groups below the minimum return None; the asynchronous collector then takes another ready prompt group. Batch-relative PPO and REINFORCE retain a usable singleton.

8. Each v1 arun_episode call is one logical rollout. Context compaction and agent.export_style: individual can export multiple rows without aborting the group. The collector records contiguous row counts and optional reward references in RolloutGroup, under the rollout_group trajectory key. Batch concatenation moves this metadata into TrajBatchMeta; splitting restores it to trajectories. Usable group sizes count logical rollouts, including for n_samples: 1.

9. Reward normalization uses one reference per logical rollout, for both group and batch statistics. When a rollout’s row rewards differ, the workflow must supply a finite rollout_reward scalar in its tensor dictionary, or on an exported InteractionWithTokenLogpReward. All supplied references within a rollout must agree. If omitted, equal row rewards provide the reference. The same centering and scaling apply to every row’s own reward; leave-one-out excludes the entire logical rollout from its reference baseline. When the computed reference standard deviation is at or below normalization epsilon, centering is retained with a divisor of 1. The built-in v1 agent workflow explicitly supplies its terminal reward for individual exports, preserving discounted row rewards and any reference already supplied. Custom workflows must declare their own reference for differing row rewards; no terminal, sum, or mean score is inferred.

10. An explicit reference is unchanged by the built-in row-length overlong penalty. Actor reward bias, scaling, and clipping apply to both rows and references. Without an explicit reference, penalized row rewards must still agree. Advantage normalization keeps its existing masked token statistics and per-token leave-one-out behavior; logical counts only select singleton fallbacks. GAE still runs separately on each row.

PPO-family actor loss remains globally token-weighted by default. Consequently, a partial group with more valid response tokens contributes more loss weight than a smaller or shorter group. This is the existing backward-compatible estimator, not an implicit claim of equal prompt weighting.

Grouped rollouts export target_slot_count, usable_slot_count, trainable_slot_count, fully_masked_group, singleton_slot_group, pre_filter_usable_slot_yield, and pre_filter_trainable_slot_yield. These count the original rollout calls; final accepted and rejected counts remain collector metrics after should_accept_fn runs. PPO training separately reports the logical usable group-size and valid-token loss-weight distributions, including per-size group_loss_weight_size_<N> metrics.

Output Shape

With group_size=4, a workflow returning [1, seq_len] tensors, and all four slots usable, the grouped output has shape [4, seq_len]. An incomplete accepted group uses its actual usable count as the leading dimension.

Implementation

From areal/infra/remote_inf_engine.py:

class GroupedRolloutWorkflow(RolloutWorkflow):
    async def arun_episode(self, engine, data):
        # Run N times concurrently
        results = await asyncio.gather(
            *[self.workflow.arun_episode(engine, data)
              for _ in range(self.group_size)]
        )

        # A normal result is usable; None is unusable.
        valid_results = [r for r in results if r is not None]
        if len(valid_results) < self.min_usable_group_size:
            return None

        # Merge based on result type
        if all_interaction_dicts(valid_results):
            return merge_dicts(valid_results)
        else:
            return concat_padded_tensors(valid_results)

Implementing Custom Workflows

To create a custom workflow:

1. Subclass 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. Prepare input
        input_ids = self.tokenizer.encode(data["prompt"])

        # 2. Generate completion
        req = ModelRequest(
            rid=uuid.uuid4().hex,
            input_ids=input_ids,
            gconfig=self.gconfig,
            tokenizer=self.tokenizer,
        )
        resp = await engine.agenerate(req)

        # 3. Compute reward
        reward = self.compute_reward(resp, data)

        # 4. Return tensor dict (or None to reject)
        if reward < 0:
            return None

        return self.build_tensor_dict(resp, reward)

2. Register with trainer:

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

Workflow Resolution

Workflows can be specified in multiple ways:

Format	Example	Description
Instance	MyWorkflow(...)	Pre-instantiated workflow
Class	MyWorkflow	Class (requires kwargs)
String path	"my_module.MyWorkflow"	Dynamic import
Agent workflow	Any class with async def run()	Wrapped with proxy support

The training system automatically resolves these to RolloutWorkflow instances.

See Also

• Agentic RL Tutorial - Training with agent frameworks

• Adding Custom Workflows - Step-by-step guide