🤖 RLHF / RLAIF Fine-Tuning with TruLens Metrics & TRL¶
Reinforcement Learning from Human/AI Feedback (RLHF / RLAIF) optimizes language model outputs by training policy networks against scalar reward functions.
TruLens feedback metrics (such as Groundedness, Relevance, and Safety) are natural reward signals for policy gradient trainers (e.g. Hugging Face TRL's GRPOTrainer and PPOTrainer).
Objective¶
In this notebook, we demonstrate how to:
- Wrap actual TruLens feedback functions into RL reward signals using
TRLRewardAdapterandRewardFunction. - Apply score transformations (mapping $[0, 1]$ feedback scores into $[-1, 1]$ symmetric rewards for policy gradients).
- Understand TRL data flow (string vs token ID decoding).
- Connect TruLens reward functions directly into a TRL training loop (
GRPOTrainer/PPOTrainer).
🛠️ Step 1: Initialize TruLens Provider & TRLRewardAdapter¶
TruLens provider feedback functions (e.g., OpenAI().relevance or OpenAI().relevance_with_cot_reasons) accept prompt and response as keyword arguments.
TRLRewardAdapter inspects the function signature at initialization and wraps it into a TRL-compatible reward function.
In [ ]:
Copied!
import os
from trulens.apps.rl import TRLRewardAdapter, RewardFunction
from trulens.providers.openai import OpenAI
# 1. Instantiate a TruLens provider feedback function
provider = OpenAI()
feedback_fn = provider.relevance # Takes (prompt: str, response: str) -> float
# 2. Create TRL Reward Adapter with 2x - 1 scaling ([0, 1] -> [-1, 1])
# Symmetric [-1, 1] rewards are recommended for policy gradient algorithms (PPO, GRPO)
reward_adapter = TRLRewardAdapter(
feedback_fn=feedback_fn,
transform="2x-1"
)
print("TRLRewardAdapter initialized successfully with TruLens OpenAI provider.")
import os
from trulens.apps.rl import TRLRewardAdapter, RewardFunction
from trulens.providers.openai import OpenAI
# 1. Instantiate a TruLens provider feedback function
provider = OpenAI()
feedback_fn = provider.relevance # Takes (prompt: str, response: str) -> float
# 2. Create TRL Reward Adapter with 2x - 1 scaling ([0, 1] -> [-1, 1])
# Symmetric [-1, 1] rewards are recommended for policy gradient algorithms (PPO, GRPO)
reward_adapter = TRLRewardAdapter(
feedback_fn=feedback_fn,
transform="2x-1"
)
print("TRLRewardAdapter initialized successfully with TruLens OpenAI provider.")
🧪 Step 2: Batch Reward Evaluation & TRL Data Flow¶
TRL Data Flow (Strings vs Token IDs):¶
- Decoded Text Strings: High-level TRL trainers (
GRPOTrainer) pass generated completions as decoded string lists (prompts: list[str],completions: list[str]) directly to reward functions. - Token ID Tensors: If working with low-level sequence generation, decode token sequences before passing them to the reward adapter using:
completions = tokenizer.batch_decode(completion_ids, skip_special_tokens=True) rewards = reward_adapter(prompts=prompts, completions=completions)
In [ ]:
Copied!
prompts = [
"Explain quantum computing in simple terms.",
"What is the capital of France?"
]
completions = [
"Quantum computing uses qubits and superposition to process information.",
"The capital of France is Paris."
]
# Evaluate batch rewards
rewards = reward_adapter(prompts=prompts, completions=completions)
for p, c, r in zip(prompts, completions, rewards):
print(f"Prompt: {p}\nCompletion: {c}\nReward: {r:+.2f}\n" + "-"*40)
prompts = [
"Explain quantum computing in simple terms.",
"What is the capital of France?"
]
completions = [
"Quantum computing uses qubits and superposition to process information.",
"The capital of France is Paris."
]
# Evaluate batch rewards
rewards = reward_adapter(prompts=prompts, completions=completions)
for p, c, r in zip(prompts, completions, rewards):
print(f"Prompt: {p}\nCompletion: {c}\nReward: {r:+.2f}\n" + "-"*40)
🚀 Step 3: Integrating with Hugging Face TRL GRPOTrainer¶
TRLRewardAdapter is compatible with TRL's GRPOTrainer and PPOTrainer via the reward_funcs parameter:
from trl import GRPOTrainer, GRPOConfig
# Pass reward_adapter directly into GRPOTrainer
trainer = GRPOTrainer(
model=model,
reward_funcs=[reward_adapter],
train_dataset=dataset,
args=GRPOConfig(
output_dir="./results",
per_device_train_batch_size=2,
num_generations=4,
),
)
# trainer.train()