# Implementing Custom Rollout Workflows for Agentic RL in AReaL

> Implement custom rollout workflows for agentic RL in AReaL. Subclass RolloutWorkflow and implement arun_episode to orchestrate LLM generation, reward calculation, and tensor packaging.

- Repository: [inclusionAI/areal](https://github.com/inclusionai/areal)
- Tags: how-to-guide
- Published: 2026-03-04

---

**To implement a custom rollout workflow for agentic RL in AReaL, subclass the abstract `RolloutWorkflow` class from [`areal/api/workflow_api.py`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/areal/utils/perf_tracer.py) for performance profiling

### Execution Flow

A standard rollout follows five distinct phases, as implemented in [`areal/workflow/rlvr.py`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/custom_task.py))
2. **Subclass `RolloutWorkflow`** from [`areal/api/workflow_api.py`](https://github.com/inclusionai/areal/blob/main/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:

```python

# 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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/areal/utils/perf_tracer.py), logging latency metrics and phase durations to the stats tracker automatically.