erl-trainer 0.3.2

Experiential Reinforcement Learning (ERL) — a thin wrapper on HuggingFace TRL's GRPOTrainer implementing the ERL algorithm with reflection, memory, and internalization.

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

pip install erl-trainer

Quick Start

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

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

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:

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:

@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:

@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