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

> Learn to implement SFT vs RL for agent post-training using the TRL library. This guide offers a clear comparison and implementation steps for optimizing your language models.

- Repository: [Bojie Li/ai-agent-book](https://github.com/bojieli/ai-agent-book)
- Tags: how-to-guide
- Published: 2026-07-26

---

**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`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/chapter7/cot-distillation/generate_data.py) that converts chain-of-thought distilled trajectories into SFT-ready format:

```json
{"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`](https://github.com/bojieli/ai-agent-book/blob/main/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:

```python
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:

```python
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"`):

```python
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:

```bash
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:

```python
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:

```python
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`):

```python
"""
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:

```bash
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`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/train_sft_trl.py) script outputs to `./outputs/sft_model`, which the RL script consumes via the `--sft_dir` argument.