Feature: bring your own algorithms (#1715)

* update loss interface

* refactor: rename grpo_loss to prime_rl_loss and consolidate loss interface

- Rename grpo_loss to prime_rl_loss (uses LossConfig directly)
- Move LossInputs, LossOutputs, LossFn from loss_interface.py into loss.py
- Remove loss_interface.py

Co-Authored-By: Claude Opus 4.5 <[email protected]>

* feat: add bring-your-own-loss support for custom loss functions

- Add CustomLossConfig with path and kwargs fields
- Add LossConfigType union (LossConfig | CustomLossConfig)
- Update setup_loss_fn to handle custom loss imports
- Add _import_object helper for dynamic imports
- Add test for custom loss configuration
- Add docs/bring-your-own-loss.md documentation

Co-Authored-By: Claude Opus 4.5 <[email protected]>

* feat: add bring-your-own advantage function support

- Add AdvantageInputs/AdvantageOutputs dataclasses
- Add CustomAdvantageConfig with path and kwargs
- Add setup_advantage_fn for custom advantage imports
- Refactor compute_advantages to use the new interface
- Add tests for custom advantage configuration
- Rename docs to bring-your-own-algorithms.md with both loss and advantage sections

Co-Authored-By: Claude Opus 4.5 <[email protected]>

* refactor: rename path to byo_function in custom configs

More descriptive name that ties into the "bring your own" concept.

Co-Authored-By: Claude Opus 4.5 <[email protected]>

* refactor: address PR review feedback

- Extract import_object to shared utils (deduplicate from loss.py and advantage.py)
- Add pydantic discriminator types for LossConfigType and AdvantageConfigType
- Rename byo_function to import_path (standard Python terminology)
- Rename grpo_advantage to default_advantage
- Fix train.py AttributeError when using CustomLossConfig
- Fix docs: per-example terminology, loss/advantage descriptions, config examples

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* docs: update prime_rl_loss description

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: handle custom loss metrics in training loop and restore extra=forbid on AdvantageConfig

- Guard mismatch_kl access in micro-step and step logging (custom loss may not emit it)
- Change AdvantageConfig base from BaseModel to BaseConfig to restore extra="forbid" validation

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* refactor: rename prime_rl_loss to default_loss_fn

Consistent with default_advantage naming.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* refactor: simplify loss setup logging

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* refactor: rename default_advantage to default_advantage_fn

Consistent with default_loss_fn naming.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

---------

Co-authored-by: Claude Opus 4.5 <[email protected]>

Author: samsja

docs/bring-your-own-algorithms.md

@@ -0,0 +1,140 @@
+# Bring Your Own Algorithms
+
+Prime-RL supports custom implementations for key algorithmic components, allowing you to experiment with different RL objectives and techniques.
+
+## 1. Custom Loss Functions
+
+The loss is computed **per-sequence** (per-sample). You provide a function that computes the loss for a single sequence, and the framework handles iteration and aggregation.
+
+### Interface
+
+```python
+from prime_rl.trainer.rl.loss import LossInputs, LossOutputs
+
+def my_custom_loss(inputs: LossInputs, **kwargs) -> LossOutputs:
+    ...
+```
+
+#### LossInputs
+
+```python
+@dataclass
+class LossInputs:
+    trainer_logprobs: Float[Tensor, "seq"]      # Log probs from current policy
+    inference_logprobs: Float[Tensor, "seq"]    # Log probs from reference policy
+    teacher_logprobs: Float[Tensor, "seq"] | None  # Optional teacher log probs
+    advantages: Float[Tensor, "seq"]            # Per-token advantages
+    loss_mask: Bool[Tensor, "seq"]              # Mask for valid tokens
+```
+
+#### LossOutputs
+
+```python
+@dataclass
+class LossOutputs:
+    loss: Float[Tensor, ""]         # Scalar loss for this sequence
+    metrics: dict[str, Tensor]      # Metrics to log
+```
+
+### Example: PPO Clipped Loss
+
+```python
+import torch
+from prime_rl.trainer.rl.loss import LossInputs, LossOutputs
+
+def ppo_clip_loss(inputs: LossInputs, clip_eps: float = 0.2) -> LossOutputs:
+    ratio = torch.exp(inputs.trainer_logprobs - inputs.inference_logprobs)
+    clipped_ratio = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps)
+
+    surr1 = ratio * inputs.advantages
+    surr2 = clipped_ratio * inputs.advantages
+
+    loss = -torch.min(surr1, surr2)[inputs.loss_mask].sum()
+
+    return LossOutputs(
+        loss=loss,
+        metrics={"clip_frac": (ratio != clipped_ratio)[inputs.loss_mask].float().mean()},
+    )
+```
+
+### Configuration
+
+```toml
+[loss]
+type = "custom"
+import_path = "my_module.ppo_clip_loss"
+kwargs = { clip_eps = 0.2 }
+```
+
+---
+
+## 2. Custom Advantage Functions
+
+Advantages are computed **per-example** (grouped by `rollouts_per_example`). You provide a function that computes advantages for a batch of examples.
+
+### Interface
+
+```python
+from prime_rl.orchestrator.advantage import AdvantageInputs, AdvantageOutputs
+
+def my_custom_advantage(inputs: AdvantageInputs, **kwargs) -> AdvantageOutputs:
+    ...
+```
+
+#### AdvantageInputs
+
+```python
+@dataclass
+class AdvantageInputs:
+    rewards: Float[Tensor, "num_examples rollouts_per_example"]
+    completion_lengths: Int[Tensor, "num_examples rollouts_per_example"]
+```
+
+#### AdvantageOutputs
+
+```python
+@dataclass
+class AdvantageOutputs:
+    advantages: Float[Tensor, "num_examples rollouts_per_example"]
+```
+
+### Example: Normalized Advantage
+
+```python
+import torch
+from prime_rl.orchestrator.advantage import AdvantageInputs, AdvantageOutputs
+
+def normalized_advantage(inputs: AdvantageInputs, eps: float = 1e-8) -> AdvantageOutputs:
+    """Normalize advantages to zero mean and unit variance per example."""
+    mean = inputs.rewards.mean(dim=1, keepdim=True)
+    std = inputs.rewards.std(dim=1, keepdim=True)
+    advantages = (inputs.rewards - mean) / (std + eps)
+    return AdvantageOutputs(advantages=advantages)
+```
+
+### Configuration
+
+```toml
+[advantage]
+type = "custom"
+import_path = "my_module.normalized_advantage"
+kwargs = { eps = 1e-8 }
+```
+
+---
+
+## Default Implementations
+
+If no custom function is specified:
+
+- **Loss**: Uses `default_loss_fn` (masked importance sampling with KL against the inference policy, and optional masking strategies)
+- **Advantage**: Uses `default_advantage_fn` (reward minus per-example baseline, a.k.a. DR-GRPO without std normalization)
+
+See `LossConfig` and `AdvantageConfig` for available parameters.
+
+## Tips
+
+- Your functions receive structured inputs via dataclasses with jaxtyping annotations
+- Return metrics as scalars or 1D tensors - they'll be aggregated automatically
+- Use the `loss_mask` / tensor shapes to handle variable-length sequences
+- Test your custom functions with the provided test patterns before training

src/prime_rl/orchestrator/advantage.py

@@ -1,13 +1,72 @@
+from dataclasses import dataclass
+from typing import Callable
+
 import torch
+from jaxtyping import Float, Int
+from torch import Tensor
+
+from prime_rl.orchestrator.config import AdvantageConfigType, CustomAdvantageConfig
+from prime_rl.utils.utils import import_object
+
+
+@dataclass
+class AdvantageInputs:
+    """Inputs for advantage computation."""
+
+    rewards: Float[Tensor, "num_problems rollouts_per_example"]
+    completion_lengths: Int[Tensor, "num_problems rollouts_per_example"]
+
+
+@dataclass
+class AdvantageOutputs:
+    """Outputs from advantage computation."""
+
+    advantages: Float[Tensor, "num_problems rollouts_per_example"]
+
+
+AdvantageFn = Callable[..., AdvantageOutputs]
+"""Type for an advantage function.
+
+Expected signature:
+    def my_advantage(inputs: AdvantageInputs, **kwargs) -> AdvantageOutputs:
+        ...
+"""
+
 
-from prime_rl.orchestrator.config import AdvantageConfig
+def default_advantage_fn(inputs: AdvantageInputs, length_weighted_mean: bool = False) -> AdvantageOutputs:
+    """Default GRPO advantage: reward minus per-problem baseline."""
+    if length_weighted_mean:
+        baseline = (inputs.rewards * inputs.completion_lengths).sum(
+            dim=1, keepdim=True
+        ) / inputs.completion_lengths.sum(dim=1, keepdim=True)
+    else:
+        baseline = inputs.rewards.mean(dim=1, keepdim=True)
+
+    return AdvantageOutputs(advantages=inputs.rewards - baseline)
+
+
+def setup_advantage_fn(config: AdvantageConfigType) -> AdvantageFn:
+    """Setup advantage function from config."""
+    if isinstance(config, CustomAdvantageConfig):
+        custom_fn = import_object(config.import_path)
+        kwargs = config.kwargs
+
+        def advantage_fn(inputs: AdvantageInputs) -> AdvantageOutputs:
+            return custom_fn(inputs, **kwargs)
+
+        return advantage_fn
+
+    def advantage_fn(inputs: AdvantageInputs) -> AdvantageOutputs:
+        return default_advantage_fn(inputs, length_weighted_mean=config.length_weighted_mean)
+
+    return advantage_fn
 
 
 def compute_advantages(
     rewards: list[float],
     completion_lengths: list[int],
     samples_per_problem: int,
-    advantage_config: AdvantageConfig | None,
+    advantage_config: AdvantageConfigType | None,
 ) -> list[float]:
     """
     Computes advantages from a flattened list of rewards, grouped by problem.
@@ -16,14 +75,17 @@ def compute_advantages(
         rewards: Flattened list of rewards where first `samples_per_problem` rewards are for the first problem
         completion_lengths: List of completion lengths for each reward
         samples_per_problem: Number of samples (and thus, rewards) per problem
-        advantage_config: Configuration for advantage computation
+        advantage_config: Configuration for advantage computation (AdvantageConfig or CustomAdvantageConfig)
     """
     if not advantage_config:
         return rewards
-    rewards = torch.tensor(rewards).view(-1, samples_per_problem)
-    lengths = torch.tensor(completion_lengths).view(-1, samples_per_problem)
-    if advantage_config.length_weighted_mean:
-        baseline = (rewards * lengths).sum(dim=1, keepdim=True) / lengths.sum(dim=1, keepdim=True)
-    else:
-        baseline = rewards.mean(dim=1, keepdim=True)
-    return (rewards - baseline).flatten().tolist()
+
+    advantage_fn = setup_advantage_fn(advantage_config)
+
+    inputs = AdvantageInputs(
+        rewards=torch.tensor(rewards).view(-1, samples_per_problem),
+        completion_lengths=torch.tensor(completion_lengths).view(-1, samples_per_problem),
+    )
+
+    result = advantage_fn(inputs)
+    return result.advantages.flatten().tolist()

src/prime_rl/orchestrator/config.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 from typing import Annotated, Any, Literal, TypeAlias
 
-from pydantic import AliasChoices, BaseModel, Field, model_validator
+from pydantic import AliasChoices, BaseModel, Discriminator, Field, Tag, model_validator
 
 from prime_rl.transport.config import FileSystemTransportConfig, TransportConfigType
 from prime_rl.utils.config import (
@@ -612,9 +612,36 @@ def validate_skip_verification(self):
 
 
 class AdvantageConfig(BaseConfig):
+    """Config for the default advantage."""
+
+    type: Literal["default"] = "default"
     length_weighted_mean: bool = False
 
 
+class CustomAdvantageConfig(BaseModel):
+    """Config for a custom external advantage function."""
+
+    type: Literal["custom"] = "custom"
+    import_path: Annotated[
+        str, Field(description="Import path to the advantage function (e.g., 'my_module.my_advantage')")
+    ]
+    kwargs: Annotated[
+        dict[str, Any], Field(default_factory=dict, description="Kwargs to pass to the advantage function")
+    ]
+
+
+def _advantage_config_discriminator(v: Any) -> str:
+    if isinstance(v, dict):
+        return v.get("type", "default")
+    return getattr(v, "type", "default")
+
+
+AdvantageConfigType: TypeAlias = Annotated[
+    Annotated[AdvantageConfig, Tag("default")] | Annotated[CustomAdvantageConfig, Tag("custom")],
+    Discriminator(_advantage_config_discriminator),
+]
+
+
 class FileSystemWeightBroadcastConfig(BaseModel):
     """Configures the filesystem weight broadcast."""
 
@@ -683,7 +710,7 @@ class OrchestratorConfig(BaseSettings):
     buffer: BufferConfig = BufferConfig()
 
     # The advantage configuration
-    advantage: AdvantageConfig | None = AdvantageConfig()
+    advantage: AdvantageConfigType | None = AdvantageConfig()
 
     # The logging configuration
     log: LogConfig = LogConfig()

src/prime_rl/trainer/rl/config.py

@@ -1,7 +1,7 @@
 from pathlib import Path
-from typing import Annotated, Literal, TypeAlias
+from typing import Annotated, Any, Literal, TypeAlias
 
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Discriminator, Field, Tag, model_validator
 
 from prime_rl.trainer.config import (
     AdamWConfig,
@@ -19,8 +19,9 @@
 
 
 class LossConfig(BaseConfig):
-    """Base config for loss."""
+    """Config for the default loss."""
 
+    type: Literal["default"] = "default"
     ratio_type: Annotated[Literal["token", "sequence"], Field(description="Type of importance ratio to use.")] = "token"
 
     token_mask_high: Annotated[
@@ -72,6 +73,26 @@ def validate_mask_bounds(self):
         return self
 
 
+class CustomLossConfig(BaseModel):
+    """Config for a custom external loss function."""
+
+    type: Literal["custom"] = "custom"
+    import_path: Annotated[str, Field(description="Import path to the loss function (e.g., 'my_module.my_loss')")]
+    kwargs: Annotated[dict[str, Any], Field(default_factory=dict, description="Kwargs to pass to the loss function")]
+
+
+def _loss_config_discriminator(v: Any) -> str:
+    if isinstance(v, dict):
+        return v.get("type", "default")
+    return getattr(v, "type", "default")
+
+
+LossConfigType: TypeAlias = Annotated[
+    Annotated[LossConfig, Tag("default")] | Annotated[CustomLossConfig, Tag("custom")],
+    Discriminator(_loss_config_discriminator),
+]
+
+
 class FakeDataLoaderConfig(BaseConfig):
     """Configures a fake data loader sampling random micro batches for debugging."""
 
@@ -130,7 +151,7 @@ class RLTrainerConfig(BaseSettings):
     data: DataLoaderConfig = DataLoaderConfig()
 
     # The loss configuration
-    loss: LossConfig = LossConfig()
+    loss: LossConfigType = LossConfig()
 
     # The optimizer configuration
     optim: Annotated[OptimizerConfigType, Field(discriminator="type")] = AdamWConfig()