Implementing Custom Rollout Workflows for Agentic RL in AReaL

To implement a custom rollout workflow for agentic RL in AReaL, subclass the abstract RolloutWorkflow class from areal/api/workflow_api.py and implement the asynchronous arun_episode method to orchestrate tokenization, LLM generation, reward calculation, and tensor packaging.

AReaL (inclusionai/areal) separates RL-from-Human-Feedback (RLHF) into modular components, with the rollout workflow serving as the critical bridge between inference and training. By implementing custom rollout workflows, you define exactly how raw prompts transform into training-ready tensor bundles containing inputs, masks, log-probabilities, and rewards.

Understanding the Rollout Workflow Architecture

The architecture centers on the abstract RolloutWorkflow API defined in areal/api/workflow_api.py. This contract requires implementing an asynchronous method that coordinates the entire episode generation pipeline from input data to batched tensors.

Core Components

The workflow interacts with several key interfaces:

• RolloutWorkflow – The abstract base class requiring async def arun_episode(self, engine: InferenceEngine, data: dict[str, Any]) -> dict[str, torch.Tensor]
• InferenceEngine – Provides agenerate(request) to produce ModelResponse objects from ModelRequest inputs
• ModelRequest / ModelResponse – Structured containers for token IDs, logits, log-probabilities, versioning metadata, and optional vision data
• AsyncRewardWrapper – Wraps synchronous reward callables for non-blocking async execution
• Tracing Utilities – @trace_session and session_context() decorators from areal/utils/perf_tracer.py for performance profiling

Execution Flow

A standard rollout follows five distinct phases, as implemented in areal/workflow/rlvr.py:

1. Input preparation – Convert raw data into input_ids using configurable tokenization
2. Request construction – Build ModelRequest with generation hyperparameters (gconfig) and tokenizer metadata
3. Generation – Execute engine.agenerate(req) inside a traced generate phase
4. Reward computation – Decode tokens and invoke the async reward function
5. Tensor packaging – Assemble batched tensors including input_ids, loss_mask, logprobs, versions, attention_mask, and rewards

Step-by-Step Implementation Guide

Follow these steps to create a production-ready custom workflow:

1. Create a module under areal/workflow/ (e.g., custom_task.py)
2. Subclass RolloutWorkflow from areal/api/workflow_api.py
3. Implement arun_episode with the signature async def arun_episode(self, engine: InferenceEngine, data: dict) -> dict[str, torch.Tensor]
4. Factor helper methods using @trace_session("reward") for reward calculation and @session_context() for generation phases
5. Return batched tensors with an extra batch dimension (unsqueeze(0)) for trainer compatibility

Minimal Implementation Example

Here is a complete custom workflow implementation:

# areal/workflow/custom_task.py

import uuid
import torch
from transformers import PreTrainedTokenizerFast
from areal.api.workflow_api import RolloutWorkflow
from areal.api.engine_api import InferenceEngine
from areal.api.io_struct import ModelRequest
from areal.utils.logging import getLogger
from areal.utils.perf_tracer import trace_session, session_context

logger = getLogger("CustomTaskWorkflow")

class CustomTaskWorkflow(RolloutWorkflow):
    """Single-generation workflow with scalar reward computation."""

    def __init__(self, reward_fn, tokenizer: PreTrainedTokenizerFast, gconfig):
        self.reward_fn = reward_fn
        self.async_reward_fn = AsyncRewardWrapper(reward_fn)
        self.tokenizer = tokenizer
        self.gconfig = gconfig.new_with_stop_and_pad_token_ids(tokenizer)

    @trace_session("reward")
    async def _compute_reward(self, resp, prompt_str, task_data):
        completions = self.tokenizer.decode(resp.output_tokens)
        return await self.async_reward_fn(prompt_str, completions, **task_data)

    @session_context()
    async def _generate(self, engine: InferenceEngine, req: ModelRequest, prompt_str, task_data):
        resp = await engine.agenerate(req)
        reward = await self._compute_reward(resp, prompt_str, task_data)
        return resp, reward

    async def arun_episode(self, engine: InferenceEngine, data: dict):
        # Tokenize prompt

        input_ids = self.tokenizer.apply_chat_template(
            data["messages"], tokenize=True, add_generation_prompt=True
        )
        req = ModelRequest(
            rid=uuid.uuid4().hex,
            input_ids=list(input_ids),
            gconfig=self.gconfig.new(n_samples=1),
            tokenizer=self.tokenizer,
        )
        prompt_str = self.tokenizer.decode(input_ids)

        # Generate and compute reward

        resp, reward = await self._generate(engine, req, prompt_str, data)

        # Build output tensors

        seq = resp.input_tokens + resp.output_tokens
        loss_mask = [0] * resp.input_len + [1] * resp.output_len
        logprobs = [0.0] * resp.input_len + resp.output_logprobs
        versions = [-1] * resp.input_len + resp.output_versions

        result = {
            "input_ids": torch.tensor(seq, dtype=torch.int32),
            "loss_mask": torch.tensor(loss_mask, dtype=torch.int32),
            "logprobs": torch.tensor(logprobs, dtype=torch.float32),
            "versions": torch.tensor(versions, dtype=torch.int32),
            "rewards": torch.tensor(reward, dtype=torch.float32),
            "attention_mask": torch.ones(len(seq), dtype=torch.bool),
        }
        # Add batch dimension for trainer compatibility

        return {k: v.unsqueeze(0) for k, v in result.items()}

Critical implementation details:

• All return tensors must include a batch dimension of 1 using unsqueeze(0)
• Use uuid.uuid4().hex for unique request IDs in ModelRequest
• Wrap synchronous reward functions with AsyncRewardWrapper to prevent blocking
• Maintain the exact tensor schema: input_ids, loss_mask, logprobs, versions, rewards, attention_mask

Implementation Patterns from Built-in Workflows

AReaL provides three reference implementations demonstrating different agentic RL patterns:

Single-Turn RL with RLVRWorkflow

areal/workflow/rlvr.py implements the canonical text-only RL-from-Verification workflow. It demonstrates standard tokenization, single-pass generation, and straightforward reward attribution. Use this as the baseline for text-based tasks with immediate reward signals.

Multi-Attempt Logic with MultiTurnWorkflow

areal/workflow/multi_turn.py extends the base pattern with retry loops. The workflow continues generating attempts until achieving a positive reward or hitting attempt limits. This pattern is essential for agentic tasks requiring iterative refinement or tool use with success verification.

Vision-Enabled Rollouts with VisionRLVRWorkflow

areal/workflow/vision_rlvr.py shows how to incorporate multimodal inputs. It utilizes AutoProcessor for image preprocessing and populates the multi_modal_input field in ModelRequest. This workflow handles vision-language tasks where rewards depend on generated descriptions of visual content.

Summary

• Subclass RolloutWorkflow from areal/api/workflow_api.py and implement arun_episode to define custom agentic RL logic
• Return batched tensors with shape (1, seq_len) using unsqueeze(0) for compatibility with AReaL's training loop
• Orchestrate four phases: input tokenization, engine.agenerate() execution, async reward computation, and tensor packaging
• Leverage tracing decorators (@trace_session, @session_context) from areal/utils/perf_tracer.py to maintain observability standards
• Reference built-in workflows (RLVRWorkflow, MultiTurnWorkflow, VisionRLVRWorkflow) for single-turn, multi-attempt, and vision-enabled patterns

Frequently Asked Questions

What tensor schema must arun_episode return for AReaL training?

The method must return a dictionary containing six specific tensors: input_ids (int32), loss_mask (int32), logprobs (float32), versions (int32), rewards (float32), and attention_mask (bool). Each tensor must include a batch dimension of 1, achieved by calling unsqueeze(0) on the sequence-level tensors before returning.

How do I handle synchronous reward functions in the async workflow?

Wrap synchronous reward functions using AsyncRewardWrapper before assignment to self.async_reward_fn. This wrapper ensures the reward computation runs in a separate thread pool, preventing blocking of the async event loop during generation batches.

Can custom workflows support multimodal inputs like images?

Yes. Following the pattern in areal/workflow/vision_rlvr.py, import AutoProcessor alongside your tokenizer, preprocess images through the processor, and pass the resulting tensors via the multi_modal_input parameter when constructing ModelRequest objects.

Where should I place tracing decorators for optimal profiling?

Apply @session_context() to generation wrapper methods and @trace_session("reward") to reward calculation methods. These decorators integrate with AReaL's performance tracer in areal/utils/perf_tracer.py, logging latency metrics and phase durations to the stats tracker automatically.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}