documents/reward_function.md

# Reward Function & Training Loop — Agentic RAG Gym

## Reward Function Design

### Philosophy

The reward function follows three core principles:
1. **Process Supervision** — Reward intermediate steps, not just final outcomes
2. **Anti-Hacking** — Detect and penalize degenerate strategies
3. **Composite Signals** — Multiple evaluation dimensions prevent reward hacking

### Composite Reward Architecture

The `CompositeRewardFunction` combines five weighted signals:

| Signal | Weight | Description |
|---|---|---|
| Retrieval Relevance | 0.25 | Average relevance score of retrieved documents |
| Reasoning Quality | 0.20 | Depth and logic of reasoning traces |
| Answer Completeness | 0.30 | Coverage of task requirements in final answer |
| Efficiency | 0.15 | Penalty for excessive steps |
| Anti-Hack Penalty | 0.10 | Deductions for detected hacking patterns |

### Step-Level Rewards (Process Supervision)

Each action type receives a specialized evaluation:

#### Retrieval Steps
- Base reward proportional to average document relevance scores
- Bonus for retrieving diverse sources
- No reward if no documents found

#### Reasoning Steps
- Trace length (>50 chars baseline)
- Evidence markers: "because", "therefore", "based on"
- Critical thinking markers: "however", "alternatively", "caveat"
- Scaled by step position (earlier = higher efficiency multiplier)

#### Answer Steps
- Length heuristic (word count thresholds)
- Source grounding: overlap between answer terms and retrieved document terms
- Rubric alignment: presence of task-specific criteria keywords

### Anti-Reward-Hacking Measures

Four detection mechanisms:

1. **Repetition Detection** — If last 3 queries are identical → 0.3 penalty
2. **Monotonic Action Exploit** — If all steps use same action type → 0.3 penalty
3. **Copy-Paste Detection** — If answer equals a previous query → 0.4 penalty
4. **Degenerate Output** — If unique word ratio < 0.3 → 0.3 penalty

### Score Clamping

All scores are strictly clamped to `[0.01, 0.99]` to ensure:
- No task returns exactly 0.0 (which would indicate a broken grader)
- No task returns exactly 1.0 (which would indicate a trivial task)

### LLM Judge Mode

An alternative `LLMJudgeRewardFunction` uses the LLM itself as an evaluator:
- Per-step: Prompts the LLM to rate step quality 0.0-1.0
- Per-episode: Prompts the LLM to rate overall performance

This is used when rule-based evaluation is insufficient for the domain.

## Training Loop

### Episode Structure

```
reset(task_id)
    │
    for step in range(max_steps):
    │   │
    │   ├── Agent observes state
    │   ├── Agent selects action (retrieve/reason/answer/plan/critique/verify)
    │   ├── Environment processes action
    │   ├── Step reward computed (process supervision)
    │   ├── State updated
    │   │
    │   └── if done: break
    │
    ├── Episode reward computed
    ├── Grader evaluates final performance
    └── Trajectory saved for training
```

### Self-Improvement Loop

The environment supports self-improvement through:

1. **Adversarial Critique** — The Critic agent identifies weaknesses in reasoning
2. **Iterative Refinement** — Retriever can be redirected based on critique
3. **Verification Gate** — Verifier checks answer grounding before submission
4. **Curriculum Difficulty** — Tasks range easy → hard, challenging frontier models

### Trajectory Analysis

Each trajectory records:
- Per-step actions and reasoning traces
- Per-step intermediate rewards
- Final score and episode metadata
- Agent communication history

This data enables:
- Offline RL training on collected trajectories
- Process reward model training
- Failure mode analysis
- Self-improvement curriculum generation

## Grading System

### Deterministic Graders

Each task has a `KeywordCoverageGrader` that:
1. Checks keyword presence across rubric categories
2. Weights coverage by category importance
3. Adds process quality bonus (20% weight)
4. Clamps final score to [0.01, 0.99]

### Process Quality Evaluation

The process component (20% of final grade) rewards:
- Diverse action types (3+ types → bonus)
- Proper sequencing (retrieve before answer)
- Reasoning steps present
- Not using excessive steps