# 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, swap `GRPOTrainer` for `ERLTrainer`, add a `feedback_func`, and you get ERL training. 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")

# Standard reward function (same as TRL)
def reward_func(prompts, completions, **kwargs):
    return [compute_your_score(c) for c in completions]

# NEW: Textual feedback function (unique to ERL)
def feedback_func(prompts, completions, **kwargs):
    return [get_your_feedback(c) for c in completions]

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,
    feedback_func=feedback_func,   # NEW: the only addition 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 textual feedback `f1`. These are two separate signals — `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.

## 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. |
| `enable_memory` | `True` | Toggle cross-episode memory on/off. |
| `enable_internalization` | `True` | Toggle the distillation step on/off. |

## Feedback Function

The only new concept vs TRL is `feedback_func`. It receives the same keyword arguments as a reward function and must return a list of strings — one per completion:

```python
def feedback_func(prompts, completions, **kwargs) -> list[str]:
    feedbacks = []
    for prompt, completion in zip(prompts, completions):
        feedbacks.append(f"Your answer was missing: {diagnose(prompt, completion)}")
    return feedbacks
```

`feedback_func` may be `None`. In that case, empty strings are passed to the reflection prompt — training still runs, but reflection quality will be lower since the model has no textual diagnosis to work from.

## Implementation Notes

### TRL version compatibility

The trainer accesses several TRL-internal methods (`_generate`, `_calculate_rewards`, `_get_per_token_logps_and_entropies`). Compatibility guards are in place for the most likely breaking changes:

- `_generate` return values are accessed by index rather than by position destructuring, so additions to TRL's return tuple are silently ignored.
- `_get_per_token_logps_and_entropies` output is handled whether TRL returns a `(logps, entropies)` tuple or a dict.
- The input batch is normalised at entry so both `dict[str, list]` (standard PyTorch DataLoader collate) and `list[dict]` (TRL's collator) are accepted.

Tested against `trl>=0.15.0`. Pin your TRL version in production.

### 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 1** (simplified) from the paper: `y1`, `Δ`, and `y2` are packed into one combined batch and updated in a single GRPO pass with global advantage normalisation.

The paper's Appendix A describes **Algorithm 2** (full), which runs two separate RL updates — one on `y1` alone, then one on `Δ + y2` — giving per-group advantage normalisation for each update. The directional gradients are the same; only the advantage scale differs. Algorithm 2 can be implemented by splitting the packed batch and calling `compute_loss` twice.

### Ablations

Both memory and internalization can be disabled independently, which is useful for ablation studies:

```python
config = ERLConfig(
    enable_memory=False,          # no cross-episode memory
    enable_internalization=False, # no distillation step, pure RL
    ...
)
```

## License

MIT
