Implementation:OpenRLHF OpenRLHF ProcessRewardDataset init

Knowledge Sources	OpenRLHF Process Reward Models
Domains	Data_Processing, Reward_Modeling
Last Updated	2026-02-07 10:40 GMT

Overview

Concrete tool for constructing step-level labeled datasets for process reward model training.

Description

The ProcessRewardDataset class extends PyTorch Dataset to handle data for process reward models. It tokenizes input sequences and assigns step-level labels at placeholder token positions. Labels can be either string tokens (e.g., "+", "-") converted to token IDs, or float reward values. The class handles truncation of both input tokens and labels, ensuring alignment between placeholder positions and reward annotations. The custom collate_fn right-pads sequences for batching.

Usage

Use this dataset class when training a process reward model (PRM) that evaluates reasoning quality at each step. The input data must contain sequences with placeholder tokens marking step boundaries and corresponding step-level labels.

Code Reference

Source Location

• Repository: OpenRLHF
• File: openrlhf/datasets/process_reward_dataset.py
• Lines: 1-107

Signature

class ProcessRewardDataset(Dataset):
    def __init__(
        self,
        dataset,
        tokenizer: Callable,
        max_length: int,
        strategy,
        multiple_of: int = 1,
    ) -> None: ...

    def __len__(self) -> int: ...
    def __getitem__(self, idx) -> Tuple[Tensor, Tensor, Tensor]: ...
    def collate_fn(self, item_list) -> Tuple[Tensor, Tensor, Tensor]: ...

Import

from openrlhf.datasets.process_reward_dataset import ProcessRewardDataset

I/O Contract

Inputs

Name	Type	Required	Description
dataset	HF Dataset	Yes	Must have columns matching input_key and label_key from strategy.args
tokenizer	Callable	Yes	HuggingFace tokenizer
max_length	int	Yes	Maximum sequence length for tokenization
strategy	DeepspeedStrategy	Yes	Provides args (input_key, label_key, placeholder_token, reward_tokens)

Outputs (__getitem__)

Name	Type	Description
input_ids	Tensor	Tokenized input sequence (1, seq_len)
attention_mask	Tensor	Attention mask (1, seq_len)
labels	Tensor	Step-level labels at placeholder positions, -100 elsewhere (1, seq_len)

Usage Examples

Creating PRM Dataset

from openrlhf.datasets.process_reward_dataset import ProcessRewardDataset
from openrlhf.datasets.utils import blending_datasets

# Load raw data
train_data = blending_datasets(
    args.dataset,
    args.dataset_probs,
    strategy,
    args.seed,
    max_count=args.max_samples,
)

# Create PRM dataset
# Data format: input contains text with placeholder tokens at step boundaries
# Labels are lists of "+" / "-" strings or float reward values
train_dataset = ProcessRewardDataset(
    train_data,
    tokenizer,
    args.max_len,
    strategy,
)

# Use with DataLoader
train_dataloader = strategy.setup_dataloader(
    train_dataset,
    args.micro_train_batch_size,
    True,
    True,
    train_dataset.collate_fn,
)

Related Pages

• Environment:OpenRLHF_OpenRLHF_CUDA_GPU_Environment