erl-trainer 0.1.4

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

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.

Compatibility

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

erl-trainer accesses several TRL-internal methods (_generate, _calculate_rewards, _get_per_token_logps_and_entropies). Two runtime guards protect against breaking changes:

• _unpack_generate — extracts prompt IDs, completion IDs, and completion texts from _generate's return value using type-based detection rather than positional indices, so additions or reordering in TRL's return signature are silently tolerated.
• _discover_batch_keys — on the first training step, pattern-matches the packed batch dict to confirm the key names that _compute_loss expects. If TRL renames a key (e.g. prompt_ids → prompt_input_ids), a warning is emitted and the correct name is used automatically.

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