erl-trainer 0.2.0

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, 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

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

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

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

erl-trainer 0.2.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.

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.

Advantage update

Because TRL 0.17.0's _compute_loss reads advantages directly from the returned dict, ERL injects its corrected advantages at the dict level before returning:

• Non-gated samples keep their y1 reward in the advantage computation.
• Gated samples have their reward replaced with the y2 reward (the second-attempt score).

The combined reward tensor is then group-normalised using the same formula as the parent, so the advantage scale is consistent across the batch.

Algorithm 1 vs Algorithm 2

This implementation follows Algorithm 1 (simplified) from the paper: advantages are derived from a combined reward tensor (y1 rewards for non-gated, y2 rewards for gated) and the GRPO update runs once over the y1 completions with those corrected advantages.

The paper's Appendix A describes Algorithm 2 (full), which runs two separate RL updates — one on y1 alone, then one on Δ + y2. Algorithm 2 can be approximated by calling compute_loss twice with different inputs.

Compatibility

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

When a new TRL minor version is released, we verify compatibility and update the version constraint.

Ablations

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

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

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