# Add Reward Function

> Guide for adding a custom reward function in vime and wiring it through --custom-rm-path (and optional reward post-processing). Use when user wants new reward logic, remote/service reward integration, or task-specific reward shaping.

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

---


# Add Reward Function

Implement custom reward logic and connect it to vime rollout/training safely.

## When to Use

Use this skill when:

- User asks to add new reward computation logic
- User asks to integrate an external reward service
- User asks to customize reward normalization/post-processing

## Step-by-Step Guide

### Step 1: Choose Reward Mode

Pick one of these:

- Single-sample mode (`--group-rm` disabled): custom function gets one `Sample`
- Group/batch mode (`--group-rm` enabled): custom function gets `list[Sample]`

`vime.rollout.rm_hub.__init__.py` calls your function via `--custom-rm-path`.

### Step 2: Create Reward Module

Create `vime/rollout/rm_hub/<your_rm>.py`.

Supported signatures:

```python
async def custom_rm(args, sample):
    return float_reward_or_reward_dict
```

```python
async def custom_rm(args, samples):
    return list_of_rewards
```

If using group mode, return one reward per sample in input order.

### Step 3: Keep Reward Type Consistent

- Return scalar numeric rewards unless your pipeline explicitly uses keyed rewards.
- If using reward dicts, ensure downstream `reward_key` / `eval_reward_key` is configured.
- Keep exceptions explicit for invalid metadata instead of silently returning zeros.

### Step 4: Optional Reward Post-Processing

To customize normalization/shaping before advantage computation, add:

```python
def post_process_rewards(args, samples):
    # return (raw_rewards, processed_rewards)
    ...
```

Wire with:

```bash
--custom-reward-post-process-path <module>.post_process_rewards
```

This hook is consumed in `vime/ray/rollout.py`.

### Step 5: Wire and Validate

Use:

```bash
--custom-rm-path vime.rollout.rm_hub.<your_rm>.custom_rm
```

## Common Mistakes

- Returning wrong output shape in group mode
- Mixing scalar rewards and reward dicts without `reward_key` config
- Doing blocking network calls without async handling
- Forgetting to validate reward behavior on truncated/failed samples

## Reference Locations

- Reward dispatch: `vime/rollout/rm_hub/__init__.py`
- Reward post-process hook: `vime/ray/rollout.py`
- Customization docs: `docs/en/get_started/customization.md`