# erl-trainer

**Experiential Reinforcement Learning (ERL)** — a thin wrapper on HuggingFace TRL's `GRPOTrainer` that adds a reflection-retry-internalization loop to standard GRPO training.

Install it and swap `GRPOTrainer` for `ERLTrainer`. Everything else — LoRA, quantization, datasets, reward functions — works exactly like TRL.

## Installation

```bash
pip install erl-trainer
```

## Quick Start

```python
from erl import ERLConfig, ERLTrainer
from datasets import load_dataset
from peft import LoraConfig  # optional

dataset = load_dataset("your_dataset", split="train")

# Simple (GRPO-compatible, no feedback):
def reward_func(prompts, completions, **kwargs):
    return [compute_your_score(p, c) for p, c in zip(prompts, completions)]

# With feedback (better reflections):
def reward_func(prompts, completions, **kwargs):
    results = []
    for p, c in zip(prompts, completions):
        score = compute_your_score(p, c)
        feedback = explain_your_score(p, c)  # e.g. "wrong: expected 4, got 5"
        results.append((score, feedback))
    return results

config = ERLConfig(
    output_dir="erl-output",
    num_train_epochs=3,
    learning_rate=1e-6,
    per_device_train_batch_size=4,
    num_generations=4,
    # ERL-specific params
    reward_threshold=1.0,
    memory_size=50,
    memory_top_k=3,
    internalization_coef=1.0,
)

# Optional: LoRA config (works exactly like TRL)
lora_config = LoraConfig(
    r=64,
    lora_alpha=64,
    target_modules="all-linear",
)

trainer = ERLTrainer(
    model="Qwen/Qwen2.5-3B-Instruct",
    args=config,
    train_dataset=dataset,
    reward_funcs=reward_func,  # that's it — no extra arguments vs TRL
    peft_config=lora_config,   # optional, same as TRL
)

trainer.train()
```

## The ERL Algorithm

Each training step runs seven phases:

| Phase | Description |
|-------|-------------|
| **1. First attempt** | Generate responses `y1` for the batch; compute numerical reward `r1` and (optionally) textual feedback `f1` from the reward function. `r1` drives the RL math; `f1` explains *why* the attempt failed. |
| **2. Gating** | Samples where `r1 < reward_threshold` enter the reflection loop; others are done. No wasted compute on already-successful attempts. |
| **3. Self-reflection** | For gated samples, the model reflects using all five inputs: the original prompt, `y1`, `f1`, `r1`, and relevant entries retrieved from cross-episode memory. Produces a natural-language improvement plan `Δ`. |
| **4. Second attempt** | The model generates `y2` conditioned on the original prompt and `Δ` only (not `y1` or `f1`). Reward `r2` is computed against the original task. |
| **5. Memory update** | If `r2 > threshold`, the reflection `Δ` is stored in a FIFO cross-episode memory. Future steps retrieve the most recent stored reflections to seed the reflection prompt. |
| **6. GRPO update** | Policy gradient over the combined batch — `y1` (reward `r1`), `Δ` (reward `r2`), and `y2` (reward `r2`) — in one joint update. Negative advantage pushes the model away from bad outputs; positive advantage reinforces good ones. |
| **7. Internalization** | SFT cross-entropy on `(original_prompt → y2)` pairs for successful second attempts (`r2 > 0`). Trains the model to produce the improved answer directly from `x`, without any reflection scaffold at inference time. |

Reflections and retries are generated by the **same model with the same weights** as the first attempt — no freezing, no separate model. All generation happens before the optimizer step.

## Reward Function Format

The reward function is the only concept from TRL that ERL extends. It may return either format:

**Format A — plain scores** (GRPO-compatible, no feedback):

```python
def reward_func(prompts, completions, **kwargs) -> list[float]:
    return [compute_score(p, c) for p, c in zip(prompts, completions)]
```

**Format B — (score, feedback) tuples** (richer reflections):

```python
def reward_func(prompts, completions, **kwargs) -> list[tuple[float, str]]:
    results = []
    for prompt, completion in zip(prompts, completions):
        score = compute_score(prompt, completion)
        feedback = f"Your answer was missing: {diagnose(prompt, completion)}"
        results.append((score, feedback))
    return results
```

Both formats work transparently. Format A gives empty-string feedback to the reflection prompt — training still runs, but reflection quality will be lower since the model has no textual diagnosis to work from. Format B provides a natural-language explanation that the model uses to write a better improvement plan.

The reward function is called **once** per training step for first-attempt completions (by TRL's parent). ERL reads the cached result for reward and feedback — no second call.

## Configuration

All `GRPOConfig` options are inherited. ERL adds:

| Parameter | Default | Description |
|-----------|---------|-------------|
| `reward_threshold` | `1.0` | Gating threshold τ. Samples with `r1 >= τ` skip reflection. |
| `memory_size` | `50` | Max reflections stored in cross-episode memory. |
| `memory_top_k` | `3` | Reflections retrieved per reflection prompt. |
| `reflection_system_prompt` | *(built-in)* | Template with `{prompt}`, `{attempt}`, `{feedback}`, `{reward}`, `{memory}`. |
| `retry_system_prompt` | *(built-in)* | Template with `{prompt}` and `{reflection}`. |
| `internalization_coef` | `1.0` | Weight of internalization loss relative to RL loss. |
| `erl_rl_coef` | `1.0` | Weight of Δ+y2 GRPO loss. Set to `0.0` for Algorithm 1 mode. |
| `enable_memory` | `True` | Toggle cross-episode memory on/off. |
| `enable_internalization` | `True` | Toggle the distillation step on/off. |
| `debug` | `False` | Enable detailed per-step logging of every ERL phase. |

## Implementation Notes

### TRL version compatibility

`erl-trainer 0.3.x` targets **TRL 0.17.x** exclusively.

In TRL 0.17.0 the monolithic `_generate_and_score_completions` method handles everything — tokenisation, generation, EOS masking, log-probability computation, reward evaluation, advantage normalisation, and metrics logging — and returns a plain dict consumed by `_compute_loss`. There is no separate `_generate`, `_calculate_rewards`, or `_get_per_token_logps_and_entropies`.

`ERLTrainer` overrides `_generate_and_score_completions` and delegates Phase 1 (first attempt) entirely to the parent. ERL phases 2–7 run on top of the parent's output. The returned dict has the same keys as the parent's method so `_compute_loss` works without modification.

### Reward function caching

Each callable reward function is wrapped in a transparent `_CachingRewardWrapper` at trainer initialisation. When TRL's parent calls the reward function during Phase 1, the wrapper captures the return value (splitting tuples into scores + feedback). ERL then reads the cache instead of calling the function again — halving the number of reward function calls per step.

`nn.Module`-based reward functions (model-based rewards) are **not** wrapped, since they use a different calling convention.

### Batched reflection and retry generation

Reflections and retries for all gated samples in a batch are generated in two batched `model.generate` calls (one for all reflections, one for all retries), not one call per sample. This keeps GPU utilisation high regardless of how many samples are gated.

### Algorithm 1 vs Algorithm 2

This implementation follows **Algorithm 2** from the ERL paper:

- **y1 GRPO update** — parent's standard GRPO loss using r1 advantages.
- **Δ + y2 GRPO update** — separate GRPO loss using batch-wide-normalised r2 advantages. Set `erl_rl_coef=0.0` to disable (Algorithm 1 mode).
- **Internalization** — SFT cross-entropy on `(prompt → y2)` for successful retries.

### Compatibility

| erl-trainer | TRL | transformers |
|-------------|-----|--------------|
| 0.3.x | 0.17.x | ≥ 4.50.0 |
| 0.2.x | 0.17.x | ≥ 4.50.0 |
| 0.1.x | 0.15.x | ≥ 4.50.0 |

### Ablations

Both memory and internalization can be disabled independently:

```python
config = ERLConfig(
    enable_memory=False,          # no cross-episode memory
    enable_internalization=False, # no distillation step, pure RL
    erl_rl_coef=0.0,              # Algorithm 1 mode (no Δ+y2 RL update)
    ...
)
```

## Citation

If you use `erl-trainer` in your research, please cite the original ERL paper and TRL:

**ERL paper** — the algorithm this package implements:

```bibtex
@article{shi2026erl,
  title   = {Experiential Reinforcement Learning},
  author  = {Shi, Taiwei and Chen, Sihao and Jiang, Bowen and Song, Linxin and Yang, Longqi and Zhao, Jieyu},
  journal = {arXiv preprint arXiv:2602.13949},
  year    = {2026},
  url     = {https://arxiv.org/abs/2602.13949}
}
```

**TRL** — the GRPO trainer this package extends:

```bibtex
@software{vonwerra2020trl,
  title  = {{TRL: Transformers Reinforcement Learning}},
  author = {von Werra, Leandro and Belkada, Younes and Tunstall, Lewis and
            Beeching, Edward and Thrush, Tristan and Lambert, Nathan and
            Huang, Shengyi and Rasul, Kashif and Gallouédec, Quentin},
  url    = {https://github.com/huggingface/trl},
  year   = {2020}
}
```

## License

MIT
