SFT vs RL for Agent Post-Training: Implementation Guide with TRL

Supervised Fine-Tuning (SFT) aligns language models to instruction formats using cross-entropy loss on curated datasets, while Reinforcement Learning (RL) optimizes agent behaviors through reward signals using policy-gradient methods like PPO, with both stages sharing the same LoRA-efficient architecture in the bojieli/ai-agent-book codebase.

Agent post-training typically follows a two-stage pipeline: first SFT to establish instruction-following capabilities using human-written examples, then RL to refine task-specific behaviors through environmental feedback. The bojieli/ai-agent-book repository provides production-ready implementations for SFT using the trl library, alongside architectural patterns for extending to RL phases.

Overview of the Post-Training Pipeline

Both phases share the same base model and tokenizer, creating a continuous workflow from behavior cloning to policy optimization:

Phase	Goal	Typical Data	Core Library	Main Script
Supervised Fine-Tuning (SFT)	Align model to desired instruction format	Human-written instruction/response pairs	trl.SFTTrainer	chapter8/prompt-distillation/train_sft_trl.py
Reinforcement Learning (RL)	Improve policy via reward signals	Auto-generated trajectories + reward model	trl.PPOTrainer	Reference implementation (see below)

The SFT step creates a solid "cold-start" policy, while the RL step refines it by maximizing a scalar reward computed from the agent's roll-outs.

Implementing Supervised Fine-Tuning (SFT)

The repository implements SFT through the trl library's SFTTrainer, which handles teacher-forcing cross-entropy loss over target token sequences.

Data Preparation

SFT requires JSONL files with a messages field containing structured conversations. The repository includes a helper script located in chapter7/cot-distillation/generate_data.py that converts chain-of-thought distilled trajectories into SFT-ready format:

{"messages": [
    {"role": "user", "content": "What is the capital of France?"},
    {"role": "assistant", "content": "Paris"}
]}

Loading and Processing Data

In chapter8/prompt-distillation/train_sft_trl.py, the load_jsonl_dataset() function (lines 24‑55) reads the JSONL and builds a Hugging Face Dataset object suitable for batching:

from datasets import load_dataset

def load_jsonl_dataset(file_path):
    return load_dataset("json", data_files=file_path, split="train")

Model Initialization with LoRA

The prepare_model_and_tokenizer() function (lines 59‑73) loads the base model with optional LoRA adapters to reduce GPU memory usage. The configuration uses r=32 and lora_alpha=16 targeting all linear layers:

from peft import LoraConfig, get_peft_model

peft_cfg = LoraConfig(
    r=32,
    lora_alpha=16,
    target_modules="all-linear",
    lora_dropout=0.0,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, peft_cfg)

Training Execution

The train_model() function (lines 45‑88) constructs a trl.SFTConfig and instantiates trl.SFTTrainer. Key configurations include mixed-precision bfloat16 training, epoch-based checkpointing (save_strategy="epoch"), and optional Weights & Biases logging (report_to="wandb"):

from trl import SFTTrainer, SFTConfig

training_args = SFTConfig(
    output_dir=output_dir,
    num_train_epochs=1,
    per_device_train_batch_size=4,
    bf16=True,
    save_strategy="epoch",
    report_to="wandb",
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    args=training_args,
)
trainer.train()

Command-Line Usage

Invoke the complete SFT pipeline directly:

python -m chapter8.prompt-distillation.train_sft_trl \
    --model_name_or_path meta-llama/Meta-Llama-3-20B-Instruct \
    --train_file ./data/sft_cot_distill.jsonl \
    --output_dir ./outputs/sft_model \
    --num_train_epochs 1 \
    --per_device_train_batch_size 4

Implementing Reinforcement Learning (RL)

While the repository does not ship a ready-made RL trainer, the same trl package provides PPOTrainer that follows identical architectural patterns to the SFT implementation.

Starting from the SFT Checkpoint

RL training begins from the SFT checkpoint (e.g., ./outputs/sft_model) to ensure the policy starts with coherent instruction-following capabilities before optimization:

model = AutoModelForCausalLM.from_pretrained(
    "./outputs/sft_model",
    torch_dtype=torch.bfloat16,
)

Defining the Reward Function

The reward function represents the only RL-specific component, replacing the static loss with a scalar signal. This can be binary (tool success/failure), scalar (BLEU/ROUGE), or learned:

def reward_fn(sample: dict) -> float:
    """
    Simple reward: +1 if assistant mentions 'correct', else 0.
    Replace with tool-usage score or LLM-based preference model.
    """
    assistant = sample["messages"][-1]["content"]
    return 1.0 if "correct" in assistant.lower() else 0.0

PPO Configuration and Training Loop

The minimal RL script below mirrors the SFT script's structure, using PPOConfig with algorithm-specific parameters (clip_range=0.2, gamma=0.99, ppo_epochs=4):

"""
RL fine-tuning of an agent using TRL PPOTrainer.
"""
import argparse
import torch
from pathlib import Path
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model
from trl import PPOTrainer, PPOConfig

def reward_fn(sample: dict) -> float:
    assistant = sample["messages"][-1]["content"]
    return 1.0 if "correct" in assistant.lower() else 0.0

def load_checkpoint(model_dir: str, use_lora: bool = True):
    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    model = AutoModelForCausalLM.from_pretrained(
        model_dir,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        use_cache=False,
    )
    
    if use_lora:
        peft_cfg = LoraConfig(
            r=32,
            lora_alpha=16,
            target_modules="all-linear",
            lora_dropout=0.0,
            bias="none",
            task_type="CAUSAL_LM",
        )
        model = get_peft_model(model, peft_cfg)
    return model, tokenizer

def train_ppo(model, tokenizer, output_dir: str, num_episodes: int = 200, batch_size: int = 4):
    ppo_cfg = PPOConfig(
        batch_size=batch_size,
        ppo_epochs=4,
        clip_range=0.2,
        gamma=0.99,
        learning_rate=2e-5,
        max_length=2048,
        entropy_coef=0.01,
        value_loss_coef=0.5,
    )
    
    trainer = PPOTrainer(
        model=model,
        tokenizer=tokenizer,
        config=ppo_cfg,
        reward_fn=reward_fn,
    )

    for ep in range(num_episodes):
        prompts = ["Explain why the sky is blue.", "Write a short poem about AI."] * batch_size
        outputs = trainer.generate(prompts, return_prompt=False)
        trainer.step(prompts, outputs)

        if (ep + 1) % 20 == 0 and trainer.is_master_process:
            ckpt = Path(output_dir) / f"ppo_step_{ep+1}"
            ckpt.mkdir(parents=True, exist_ok=True)
            trainer.model.save_pretrained(str(ckpt))
            tokenizer.save_pretrained(str(ckpt))
            print(f"[PPO] checkpoint saved at {ckpt}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--sft_dir", required=True)
    parser.add_argument("--output_dir", default="./outputs/rl_ppo")
    parser.add_argument("--num_episodes", type=int, default=200)
    args = parser.parse_args()

    model, tokenizer = load_checkpoint(args.sft_dir)
    train_ppo(model, tokenizer, args.output_dir, args.num_episodes)

Run the RL phase by pointing to your SFT checkpoint:

python train_rl_ppo.py --sft_dir ./outputs/sft_model --num_episodes 200

Key Implementation Differences

• SFT uses teacher-forcing with static instruction-response pairs, optimizing cross-entropy loss against human-written targets. The SFTTrainer handles batching and gradient accumulation automatically.
• RL uses policy gradients with dynamically generated roll-outs, optimizing expected reward through PPOTrainer.step(). This requires environment interaction or a reward model to compute scalar feedback.
• Memory Efficiency: Both phases use identical LoRA configurations (r=32, lora_alpha=16), allowing you to load the SFT adapter directly into the RL trainer without merging weights.

Summary

• SFT establishes baseline instruction-following via trl.SFTTrainer in chapter8/prompt-distillation/train_sft_trl.py, using JSONL data with messages format and LoRA for memory efficiency.
• RL refines behaviors via trl.PPOTrainer, starting from the SFT checkpoint and optimizing a custom reward_fn that evaluates tool use or response quality.
• Both stages share architectural patterns: bfloat16 precision, epoch-based checkpointing, and compatible LoRA adapters.
• The repository provides complete SFT implementation and reference patterns for RL extension, supporting iterative SFT→RL→SFT cycles common in modern agent research.

Frequently Asked Questions

What is the main difference between SFT and RL in agent training?

SFT aligns the model to specific output formats using supervised cross-entropy loss on fixed datasets, effectively teaching the model "what to say." RL optimizes the model's decision-making process by sampling roll-outs and rewarding desirable outcomes, teaching the model "how to act" to maximize task success. According to the bojieli/ai-agent-book source code, SFT provides the cold-start policy that RL subsequently refines.

Can I use the same LoRA configuration for both SFT and RL?

Yes. The repository uses identical LoRA settings (r=32, lora_alpha=16, targeting all linear layers) for both phases. This compatibility allows you to load the SFT checkpoint directly into the RL trainer without merging adapters or converting weights, significantly reducing GPU memory requirements during the RL phase.

How do I design a reward function for agent RL training?

The reward function should return a scalar value representing task success. For tool-using agents, return +1 for successful tool execution and 0 otherwise. For quality optimization, use automated metrics like BLEU or a learned preference model. The example reward_fn in the RL implementation checks for specific keywords, but production systems should implement environment-specific logic or LLM-based evaluators.

Should I run SFT before RL or can I skip directly to RL?

Always run SFT before RL. The SFT checkpoint provides coherent instruction-following capabilities that serve as the initial policy for RL. Starting RL from a base model without SFT results in unstable training and poor sample efficiency because the policy lacks basic language understanding. The train_sft_trl.py script outputs to ./outputs/sft_model, which the RL script consumes via the --sft_dir argument.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}