eole-nlp · Copilot · Jul 8, 2026 · Jul 8, 2026 · Jul 8, 2026 · Jul 13, 2026
diff --git a/eole/config/training.py b/eole/config/training.py
@@ -244,6 +244,46 @@ class TrainingConfig(
         default="mean", description="How to aggregate attention entropy across different attention types/layers."
     )
 
+    # RL fine-tuning configuration.
+    # Only "reinforce" is implemented for now; "dpo"/"grpo"/"ppo" are reserved
+    # names for planned follow-up work and will raise NotImplementedError.
+    rl_algorithm: Literal["reinforce", "dpo", "grpo", "ppo"] | None = Field(
+        default=None,
+        description="RL fine-tuning algorithm to use instead of standard supervised training. "
+        "Only 'reinforce' is currently implemented; 'dpo'/'grpo'/'ppo' are reserved for future support. "
+        "Leave unset (None) for standard supervised (teacher-forced) training.",
+    )
+    rl_reward_metric: str | None = Field(
+        default=None,
+        description="Name of the scorer (see valid_metrics/scorers registry) used as the reward "
+        "signal for RL training. Required when rl_algorithm is set.",
+    )
+    rl_baseline: Literal["batch_mean", "none"] = Field(
+        default="batch_mean",
+        description="Variance-reduction baseline subtracted from the reward before computing the "
+        "policy-gradient loss (REINFORCE). 'batch_mean' subtracts the mean reward of the batch.",
+    )
+    rl_kl_coef: float = Field(
+        default=0.0,
+        description="Weight of the KL-divergence penalty against a frozen reference model's "
+        "log-probs. 0 disables the penalty (and the reference model is not used).",
+    )
+    rl_reference_model: str | None = Field(
+        default=None,
+        description="Path to a frozen reference model checkpoint used for the KL penalty. "
+        "Required only when rl_kl_coef != 0.",
+    )
+    rl_num_rollouts: int = Field(
+        default=1,
+        description="Number of sampled generations (rollouts) per prompt. Not used by REINFORCE "
+        "(single rollout per prompt); reserved for group-relative methods (e.g. GRPO) which need "
+        "several rollouts per prompt to compute a group baseline. Currently ignored.",
+    )
+    rl_gen_max_length: int = Field(default=256, description="Maximum generation length for RL rollouts.")
+    rl_gen_temperature: float = Field(default=1.0, description="Sampling temperature for RL rollouts.")
+    rl_gen_top_k: int = Field(default=0, description="Top-k sampling for RL rollouts (0 disables it).")
+    rl_gen_top_p: float = Field(default=0.0, description="Top-p (nucleus) sampling for RL rollouts (0 disables it).")
+
     @computed_field
     @cached_property
     def storage_dtype(self) -> torch.dtype:
@@ -320,6 +360,17 @@ def _validate_running_config(self):
         if self.lora_embedding and (self.freeze_encoder or self.freeze_decoder):
             raise ValueError("Cannot use LoRa embedding with Enc/Dec-oder freezing")
 
+        if self.rl_algorithm is not None:
+            if self.rl_algorithm != "reinforce":
+                raise NotImplementedError(
+                    f"rl_algorithm={self.rl_algorithm!r} is planned but not implemented yet. "
+                    "Only 'reinforce' is currently supported."
+                )
+            if not self.rl_reward_metric:
+                raise ValueError("rl_reward_metric must be set when rl_algorithm is set.")
+            if self.rl_kl_coef != 0.0 and not self.rl_reference_model:
+                raise ValueError("rl_reference_model must be set when rl_kl_coef != 0.")
+
         return self
 
 

diff --git a/eole/tests/test_rl_loss.py b/eole/tests/test_rl_loss.py
@@ -0,0 +1,140 @@
+import unittest
+
+import torch
+
+from eole.config.training import TrainingConfig
+from eole.utils.rl_loss import REINFORCELossCompute, build_rl_loss_compute
+
+
+class DummyVocab(dict):
+    def __getitem__(self, key):
+        if isinstance(key, str) and not dict.__contains__(self, key):
+            return 1  # pad id
+        return dict.__getitem__(self, key)
+
+
+class TestREINFORCELossCompute(unittest.TestCase):
+    def _loss_compute(self, baseline="batch_mean", kl_coef=0.0):
+        return REINFORCELossCompute(padding_idx=1, baseline=baseline, kl_coef=kl_coef)
+
+    def test_invalid_baseline_raises(self):
+        with self.assertRaises(ValueError):
+            REINFORCELossCompute(padding_idx=1, baseline="not-a-baseline")
+
+    def test_reward_magnitude_scales_the_loss(self):
+        # With a 2-sample batch and the "batch_mean" baseline, doubling the
+        # spread between rewards should scale the resulting (non-zero) loss.
+        token_log_probs = torch.tensor([[-0.1, -0.2, -0.3], [-0.5, -0.4, 0.0]])
+        padding_mask = torch.tensor([[True, True, True], [True, True, False]])
+
+        lc = self._loss_compute()
+        loss_r1, _ = lc.compute_policy_loss(token_log_probs.clone(), padding_mask, torch.tensor([1.0, 0.0]))
+        loss_r2, _ = lc.compute_policy_loss(token_log_probs.clone(), padding_mask, torch.tensor([2.0, 0.0]))
+        self.assertAlmostEqual(loss_r2.item(), 2 * loss_r1.item(), places=5)
+
+    def test_batch_mean_baseline_zero_reward_signal_when_equal(self):
+        token_log_probs = torch.tensor([[-0.1, -0.2], [-0.3, -0.1]], requires_grad=True)
+        padding_mask = torch.tensor([[True, True], [True, True]])
+        rewards = torch.tensor([1.0, 1.0])
+
+        lc = self._loss_compute(baseline="batch_mean")
+        loss, stats = lc.compute_policy_loss(token_log_probs, padding_mask, rewards)
+        self.assertAlmostEqual(loss.item(), 0.0, places=6)
+        self.assertAlmostEqual(stats["reward_mean"], 1.0)
+
+    def test_padding_positions_are_ignored(self):
+        # 2-sample batch (matching `rewards`); the huge negative log-prob at
+        # the padded position (index 2) must not blow up the loss.
+        token_log_probs = torch.tensor([[-0.1, -0.2, -100.0], [-0.1, -0.2, -100.0]], requires_grad=True)
+        padding_mask = torch.tensor([[True, True, False], [True, True, False]])
+        rewards = torch.tensor([1.0, 0.0])
+
+        lc = self._loss_compute()
+        loss, _ = lc.compute_policy_loss(token_log_probs, padding_mask, rewards)
+        self.assertTrue(torch.isfinite(loss))
+
+    def test_kl_penalty_increases_loss_when_policies_diverge(self):
+        token_log_probs = torch.tensor([[-0.1, -0.2]], requires_grad=True)
+        ref_log_probs = torch.tensor([[-1.1, -1.2]])
+        padding_mask = torch.tensor([[True, True]])
+        rewards = torch.tensor([1.0])
+
+        lc_no_kl = self._loss_compute(kl_coef=0.0)
+        lc_with_kl = self._loss_compute(kl_coef=1.0)
+
+        loss_no_kl, stats_no_kl = lc_no_kl.compute_policy_loss(
+            token_log_probs, padding_mask, rewards, ref_token_log_probs=ref_log_probs
+        )
+        loss_with_kl, stats_with_kl = lc_with_kl.compute_policy_loss(
+            token_log_probs, padding_mask, rewards, ref_token_log_probs=ref_log_probs
+        )
+        self.assertNotIn("kl", stats_no_kl)
+        self.assertIn("kl", stats_with_kl)
+        self.assertGreater(loss_with_kl.item(), loss_no_kl.item())
+
+    def test_backward_populates_gradients(self):
+        token_log_probs = torch.tensor([[-0.1, -0.2, -0.3], [-0.5, -0.4, 0.0]], requires_grad=True)
+        padding_mask = torch.tensor([[True, True, True], [True, True, False]])
+        rewards = torch.tensor([1.0, 0.0])
+
+        lc = self._loss_compute()
+        loss, _ = lc.compute_policy_loss(token_log_probs, padding_mask, rewards)
+        loss.backward()
+        self.assertIsNotNone(token_log_probs.grad)
+        self.assertTrue(torch.any(token_log_probs.grad != 0))
+
+
+class TestBuildRLLossCompute(unittest.TestCase):
+    def test_reinforce_algorithm_builds_reinforce_loss(self):
+        vocabs = {"specials": {}, "tgt": DummyVocab()}
+        config = _FakeConfig(rl_algorithm="reinforce", rl_baseline="none", rl_kl_coef=0.0)
+        loss_compute = build_rl_loss_compute(config, vocabs)
+        self.assertIsInstance(loss_compute, REINFORCELossCompute)
+        self.assertEqual(loss_compute.baseline, "none")
+
+    def test_dpo_algorithm_raises_not_implemented(self):
+        vocabs = {"specials": {}, "tgt": DummyVocab()}
+        config = _FakeConfig(rl_algorithm="dpo", rl_baseline="batch_mean", rl_kl_coef=0.0)
+        with self.assertRaises(NotImplementedError):
+            build_rl_loss_compute(config, vocabs)
+
+
+class _FakeTraining:
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
+
+
+class _FakeConfig:
+    def __init__(self, **kwargs):
+        self.training = _FakeTraining(**kwargs)
+
+
+class TestTrainingConfigRLValidation(unittest.TestCase):
+    def test_default_disables_rl(self):
+        config = TrainingConfig()
+        self.assertIsNone(config.rl_algorithm)
+
+    def test_unimplemented_algorithm_raises_not_implemented(self):
+        with self.assertRaises(NotImplementedError):
+            TrainingConfig(rl_algorithm="dpo")
+
+    def test_reinforce_requires_reward_metric(self):
+        with self.assertRaises(ValueError):
+            TrainingConfig(rl_algorithm="reinforce")
+
+    def test_reinforce_with_kl_requires_reference_model(self):
+        with self.assertRaises(ValueError):
+            TrainingConfig(rl_algorithm="reinforce", rl_reward_metric="BLEU", rl_kl_coef=0.1)
+
+    def test_valid_reinforce_config(self):
+        config = TrainingConfig(
+            rl_algorithm="reinforce",
+            rl_reward_metric="BLEU",
+            rl_kl_coef=0.1,
+            rl_reference_model="/tmp/ref_model.pt",
+        )
+        self.assertEqual(config.rl_algorithm, "reinforce")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/eole/trainer.py b/eole/trainer.py
@@ -660,6 +660,23 @@ def build_trainer(config, device_id, model, vocabs, optim, model_saver=None):
         estim_loss_lambda_steps=running_config.estim_loss_lambda_steps,
     )
 
+    if running_config.rl_algorithm is not None:
+        return _build_rl_trainer(
+            config,
+            device_id,
+            model,
+            vocabs,
+            optim,
+            trainer_config,
+            train_loss,
+            valid_loss,
+            scoring_preparator,
+            valid_scorers,
+            report_manager,
+            model_saver,
+            earlystopper,
+        )
+
     return Trainer(
         model=model,
         train_loss=train_loss,
@@ -672,3 +689,73 @@ def build_trainer(config, device_id, model, vocabs, optim, model_saver=None):
         model_saver=model_saver,
         earlystopper=earlystopper,
     )
+
+
+def _build_rl_trainer(
+    config,
+    device_id,
+    model,
+    vocabs,
+    optim,
+    trainer_config,
+    train_loss,
+    valid_loss,
+    scoring_preparator,
+    valid_scorers,
+    report_manager,
+    model_saver,
+    earlystopper,
+):
+    """Build an :class:`eole.trainer_rl.RLTrainer` for on-policy RL fine-tuning
+    (e.g. REINFORCE). Split out from :func:`build_trainer` to keep RL-specific
+    imports/wiring (rl_loss, generate_utils, reference model loading) local to
+    the RL path, and to avoid a circular import at module load time.
+    """
+    from eole.trainer_rl import RLTrainer
+    from eole.utils.rl_loss import build_rl_loss_compute
+    from eole.utils.generate_utils import GenerationConfig
+    from eole.models.model import get_model_class, get_metadata
+
+    running_config = config.training
+
+    rl_loss_compute = build_rl_loss_compute(config, vocabs)
+
+    reward_scorers_cls = get_scorers_cls([running_config.rl_reward_metric])
+    reward_scorers = build_scorers(config, reward_scorers_cls)
+
+    reference_model = None
+    if running_config.rl_kl_coef != 0.0:
+        _ref_meta = get_metadata(running_config.rl_reference_model)
+        _ref_model_class = get_model_class(_ref_meta["config"].model)
+        reference_model, _, _ = _ref_model_class.for_inference(
+            config, device_id=device_id, model_path=running_config.rl_reference_model
+        )
+        reference_model.eval()
+
+    rl_gen_config = GenerationConfig(
+        max_length=running_config.rl_gen_max_length,
+        temperature=running_config.rl_gen_temperature,
+        top_k=running_config.rl_gen_top_k,
+        top_p=running_config.rl_gen_top_p,
+    )
+
+    return RLTrainer(
+        model=model,
+        train_loss=train_loss,
+        valid_loss=valid_loss,
+        scoring_preparator=scoring_preparator,
+        valid_scorers=valid_scorers,
+        optim=optim,
+        config=trainer_config,
+        rl_loss_compute=rl_loss_compute,
+        reward_scorers=reward_scorers,
+        rl_reward_metric=running_config.rl_reward_metric,
+        vocabs=vocabs,
+        model_config=config.model,
+        device_id=device_id,
+        rl_gen_config=rl_gen_config,
+        reference_model=reference_model,
+        report_manager=report_manager,
+        model_saver=model_saver,
+        earlystopper=earlystopper,
+    )