Skip to content

Add REINFORCE RL training support (first step toward REINFORCE, DPO, GRPO, PPO) - #400

Open
vince62s with Copilot wants to merge 4 commits into
mainfrom
copilot/add-reinforce-and-dpo
Open

Add REINFORCE RL training support (first step toward REINFORCE, DPO, GRPO, PPO)#400
vince62s with Copilot wants to merge 4 commits into
mainfrom
copilot/add-reinforce-and-dpo

Conversation

Copilot AI commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

Following the merged "RL readiness refactoring" (#389), the request was to add RL fine-tuning support for REINFORCE and DPO. Implementing both simultaneously wasn't practical: REINFORCE is on-policy and can directly reuse the generation/reward infra already added in #389, while DPO needs an offline chosen/rejected preference-pair data pipeline that doesn't exist yet. This PR implements REINFORCE first, with the abstractions designed so DPO/GRPO/PPO can follow incrementally.

RL loss

  • New eole/utils/rl_loss.py: RLLossCompute base class + REINFORCELossCompute, a score-function policy-gradient loss with a batch-mean baseline and an optional KL penalty against a frozen reference model.

RL trainer

  • New eole/trainer_rl.py: RLTrainer(Trainer) overriding only _train_step() — samples rollouts, scores them with a configured scorer, recomputes differentiable log-probs of the sampled tokens, and backprops the policy-gradient loss. Reuses the outer Trainer.train() loop (validation, checkpointing, reporting) unchanged.

Generation utilities

  • eole/utils/generate_utils.py: generate_from_batch() now also returns references rebuilt from batch['tgt'] when present; generate_and_score() uses them automatically as texts_ref when not explicitly passed.

Configuration

  • eole/config/training.py: new rl_algorithm, rl_reward_metric, rl_baseline, rl_kl_coef, rl_reference_model, rl_num_rollouts, rl_gen_* fields. dpo/grpo/ppo are reserved algorithm names that raise NotImplementedError for now, so future support doesn't require breaking config changes.
  • eole/trainer.py: build_trainer() dispatches to RLTrainer when rl_algorithm="reinforce".
training:
    rl_algorithm: reinforce
    rl_reward_metric: "BLEU"
    rl_baseline: batch_mean
    rl_kl_coef: 0.05
    rl_reference_model: base_model/checkpoint

Tests & examples

  • eole/tests/test_rl_loss.py: unit tests for the REINFORCE loss (baseline behavior, padding handling, KL penalty, gradient flow) and RL config validation.
  • recipes/rl/reinforce.yaml + recipes/rl/README.md: example config and a status table tracking REINFORCE/DPO/GRPO/PPO.

Follow-up (not in this PR)

  • DPO: needs an offline chosen/rejected preference-pair data pipeline.
  • GRPO: extend RLTrainer/rl_loss with a multi-rollout group-relative baseline (rl_num_rollouts is reserved for this, currently unused by REINFORCE).
  • PPO: needs a value model and a clipped surrogate objective.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds initial on-policy RL fine-tuning support to the training stack by introducing a REINFORCE loss, an RL-specific trainer that plugs into the existing Trainer._train_step() extension point, and small generation-utility enhancements to streamline reward scoring.

Changes:

  • Introduces RLLossCompute + REINFORCELossCompute and an RL loss factory (build_rl_loss_compute).
  • Adds RLTrainer that performs rollout generation + reward scoring + differentiable log-prob recomputation + policy-gradient backprop.
  • Extends generation utilities to return detokenized references from batch["tgt"] and auto-wire them into scoring when texts_ref is not provided.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.

Show a summary per file
File Description
recipes/rl/reinforce.yaml Adds an example REINFORCE training config (RL-specific training: fields + sampling params).
recipes/rl/README.md Documents RL recipe status and explains the REINFORCE workflow and key config fields.
eole/utils/rl_loss.py Implements REINFORCE policy-gradient loss (baseline + optional KL penalty) and a factory for RL loss selection.
eole/utils/generate_utils.py Adds optional references output and auto-uses them as texts_ref in generate_and_score().
eole/trainer.py Dispatches to an RL trainer builder when training.rl_algorithm is set and wires RL-specific components.
eole/trainer_rl.py New RLTrainer that overrides _train_step() to run rollout-based RL updates.
eole/tests/test_rl_loss.py Adds unit tests for REINFORCE loss behavior and RL config validation.
eole/config/training.py Adds RL-related config fields and validates supported algorithms/required RL settings.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread eole/trainer.py
Comment thread eole/trainer_rl.py
Comment thread eole/trainer_rl.py
Comment thread recipes/rl/reinforce.yaml Outdated
@vince62s
vince62s marked this pull request as ready for review July 17, 2026 11:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants