Implementation:OpenRLHF OpenRLHF Compute reward

Knowledge Sources	OpenRLHF
Domains	Reinforcement_Learning
Last Updated	2026-02-07 00:00 GMT

Overview

Concrete tool for shaping rewards with KL penalties for PPO training provided by OpenRLHF.

Description

The compute_reward function combines scalar environment rewards with per-token KL divergence penalties. It places the environment reward at the EOS token position (found via action_mask), adds the KL penalty ($-kl\mathrm{\_}coef\times kl$) at every token position, and optionally clips rewards to a specified range.

Usage

Called during PPO experience generation after computing KL divergence and obtaining environment rewards.

Code Reference

Source Location

• Repository: OpenRLHF
• File: openrlhf/models/utils.py
• Lines: L44-72

Signature

def compute_reward(
    r: Union[torch.Tensor, float],                # Scalar reward per sequence
    kl_coef: float,                                # KL penalty coefficient
    kl: Union[torch.Tensor, list[torch.Tensor]],   # Per-token KL estimates
    action_mask: Optional[torch.Tensor] = None,    # Token-level action mask
    reward_clip_range: Tuple[float, float] = None, # Reward clipping bounds
) -> Union[torch.Tensor, list[torch.Tensor]]:
    """Returns per-token shaped reward (batch, seq)."""

Import

from openrlhf.models.utils import compute_reward

I/O Contract

Inputs

Name	Type	Required	Description
r	Tensor or float	Yes	Scalar reward per sequence (batch_size,)
kl_coef	float	Yes	KL penalty coefficient (beta)
kl	Tensor	Yes	Per-token KL estimates (batch, seq)
action_mask	Tensor	Yes	Binary mask for action tokens
reward_clip_range	Tuple	No	(min, max) reward clipping bounds

Outputs

Name	Type	Description
reward	Tensor	Per-token shaped reward (batch, seq)

Usage Examples

from openrlhf.models.utils import compute_reward, compute_approx_kl

# Compute KL penalty
kl = compute_approx_kl(policy_log_probs, ref_log_probs)

# Shape rewards
shaped_rewards = compute_reward(
    r=rm_scores,
    kl_coef=0.1,
    kl=kl,
    action_mask=action_mask,
)

Related Pages

Implements Principle

• Principle:OpenRLHF_OpenRLHF_Reward_Shaping