# Add Rollout Function

> Guide for adding a new rollout function in vime and wiring it through --rollout-function-path. Use when user wants to implement custom rollout data generation logic, custom train/eval rollout outputs, or migrate from the default vLLM rollout path.

- Skill: `vllm-project/add-rollout-function` (Agent Skill)
- Install (CLI): `npx skillmds@latest add vllm-project/add-rollout-function`
- Raw SKILL.md: https://api.skillmd.com/api/skills/vllm-project/add-rollout-function/raw
- Safety review: pending
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: Coding & Dev Tools
- Author: vllm-project (https://skillmd.com/u/vllm-project)
- Updated: 2026-09-17
- Page: https://skillmd.com/skills/vllm-project/add-rollout-function

---


# Add Rollout Function

Implement a custom rollout function and integrate it safely with vime training/eval flow.

## When to Use

Use this skill when:

- User asks to add a new rollout task or rollout generation function
- User asks to replace default `vime.rollout.vllm_rollout.generate_rollout`
- User asks to customize train/eval data generation behavior

## Step-by-Step Guide

### Step 1: Choose the Right Starting Point

Start from one of these references:

- Async RL-style rollout: `vime/rollout/vllm_rollout.py`
- Simple SFT-style rollout: `vime/rollout/sft_rollout.py`

If the task needs engine-based async generation and rewards, use the vLLM path as base.
If the task is file/buffer-driven and simple, use sft path as base.

### Step 2: Create the New Rollout Module

Create a new file, for example: `vime/rollout/<your_rollout>.py`

Required callable signature:

```python
def generate_rollout(args, rollout_id, data_source, evaluation=False) -> RolloutFnTrainOutput | RolloutFnEvalOutput:
    ...
```

Return types are defined in `vime/rollout/base_types.py`.

### Step 3: Implement Train and Eval Branches Explicitly

- For training (`evaluation=False`), return `RolloutFnTrainOutput(samples=..., metrics=...)`
- For evaluation (`evaluation=True`), return `RolloutFnEvalOutput(data=..., metrics=...)`

Minimal skeleton:

```python
from vime.rollout.base_types import RolloutFnTrainOutput, RolloutFnEvalOutput


def generate_rollout(args, rollout_id, data_source, evaluation=False):
    if evaluation:
        result = {
            "custom_eval": {
                "rewards": [],
                "truncated": [],
                "samples": [],
            }
        }
        return RolloutFnEvalOutput(data=result)

    groups = data_source.get_samples(args.rollout_batch_size)
    # fill Sample fields needed by training: tokens/response_length/reward/status (+ loss_mask when needed)
    return RolloutFnTrainOutput(samples=groups)
```

### Step 4: Keep Data Contract Compatible

For each generated sample, ensure required training fields are populated consistently with your objective:

- `tokens`
- `response_length`
- `reward` (or reward dict if your setup uses keyed rewards)
- `status`

If partial rollout or masking logic is involved, keep `loss_mask` semantics consistent with existing behavior.

### Step 5: Wire Through Arguments

Set your function path via CLI:

```bash
--rollout-function-path vime.rollout.<your_rollout>.generate_rollout
```

The default and signature expectation are documented in:

- `vime/utils/arguments.py`
- `docs/en/get_started/customization.md`

## Common Mistakes

- Returning raw Python lists/dicts with mismatched schema in custom path
- Implementing only training branch and forgetting evaluation branch
- Generating samples without required fields (`tokens`, `response_length`, `reward`, `status`)
- Using blocking-heavy logic in high-frequency rollout paths without batching/concurrency control

## Reference Locations

- Default rollout: `vime/rollout/vllm_rollout.py`
- Simple custom example: `vime/rollout/sft_rollout.py`
- Output dataclasses: `vime/rollout/base_types.py`
- Wiring/loading: `vime/ray/rollout.py`
- Argument definition: `vime/utils/arguments.py`
- Customization docs: `docs/en/get_started/customization.md`