Reward Shaping for Language Agents: From Format Checks to Trajectory Rewards

The previous three subtopics established WHAT reward functions do (score correctness, drive learning via GRPO) but not HOW to actually design a good one — arguably the single hardest practical skill in both RFT (sft-vs-rft topic) and ART-based agent training (previous subtopic). This subtopic covers that design problem specifically for TEXT-based tasks.

This course's own GRPO walkthrough (in this same fine-tuning chapter, immediately preceding the RL-environments sections) demonstrates a genuinely useful pattern by using FOUR separate reward functions rather than one: match format exactly, match format approximately, check the answer, and check numbers. Layering multiple checks like this, rather than relying on one all-or-nothing reward, gives the model partial credit and a smoother learning signal — an output that's almost right (correct answer, slightly wrong format) scores better than a completely wrong one, which is far more informative for learning than a single binary 'correct or not' signal that can't distinguish 'close' from 'nowhere close.'

This same layering principle extends naturally to TRAJECTORY-level rewards for language agents (the ART framework's domain). A few practical categories worth knowing: outcome rewards check whether the FINAL result was correct (did the agent's answer match, did the task actually get completed) — the most important signal, but the sparsest, since it only fires once at the very end of a potentially long trajectory. Process rewards score intermediate steps along the way (did the agent call an appropriate tool for this situation, did it avoid an obviously wasteful repeated action) — denser and more frequent than outcome rewards, which helps address the classic RL problem of SPARSE rewards (if the only signal is 'did the whole 20-step trajectory succeed,' the model gets very little information about which of those 20 steps actually helped). Efficiency penalties discourage unnecessarily long trajectories or redundant tool calls — directly mirroring this course's own reward-function pattern of rewarding not just correctness but ADHERENCE to a desired shape (its 'match format' checks reward the RIGHT SHAPE of answer, not just any correct-content answer).

The overarching design principle, generalizing from this course's own deterministic, hard-to-game reward functions: prefer reward signals that are automatically checkable and resistant to gaming over ones that require subjective judgment. A reward like 'did the final numeric answer match' is nearly impossible for the model to satisfy without actually solving the problem. A vaguer reward like 'did the response seem helpful' (if scored by a weaker automated check rather than a careful LLM-as-judge, from the LLM-evaluation unit) is far easier for a model to learn to superficially satisfy without genuinely improving — exactly the reward-hacking failure mode the sft-vs-rft topic's tradeoffs subtopic warned about.

# Implementing a layered, multi-check reward function in the style
# of the book's 4 GRPO reward functions, extended to a trajectory-
# level agent task (outcome + process + efficiency rewards combined).

def format_exact_match(output: str, expected_format: str) -> float:
    return 1.0 if output.strip() == expected_format else 0.0

def format_approx_match(output: str, required_tags: list[str]) -> float:
    """Partial credit: how many of the required structural tags are present."""
    present = sum(1 for tag in required_tags if tag in output)
    return present / len(required_tags) if required_tags else 0.0

def outcome_reward(trajectory: list[dict], expected_answer: str) -> float:
    """Sparse: fires only on the final trajectory outcome."""
    final = next((s for s in trajectory if s["type"] == "final_answer"), None)
    return 1.0 if final and expected_answer in final["content"] else 0.0

def process_reward(trajectory: list[dict]) -> float:
    """Denser: rewards appropriate tool use along the way, not just the end."""
    good_tool_calls = sum(1 for s in trajectory if s["type"] == "tool_call")
    return min(good_tool_calls * 0.2, 0.6)   # capped partial credit

def efficiency_penalty(trajectory: list[dict], ideal_length: int = 4) -> float:
    """Discourage unnecessarily long/redundant trajectories."""
    excess_steps = max(0, len(trajectory) - ideal_length)
    return -0.1 * excess_steps

def combined_trajectory_reward(trajectory: list[dict], expected_answer: str) -> float:
    """Layering multiple checks (the book's own pattern) instead of
    one all-or-nothing signal -- smoother, more informative learning signal."""
    return (
        outcome_reward(trajectory, expected_answer)
        + process_reward(trajectory)
        + efficiency_penalty(trajectory)
    )

example_trajectory = [
    {"type": "reasoning", "content": "thinking..."},
    {"type": "tool_call", "tool": "search"},
    {"type": "final_answer", "content": "The capital is Paris"},
]
reward = combined_trajectory_reward(example_trajectory, expected_answer="Paris")
print(f"combined trajectory reward: {reward:.2f}")