Skip to content

Alignment Losses

xaytune supports six alignment methods. Each has a dedicated loss function and can be selected via create_alignment_loss_fn().

Method Function Paper
DPO dpo_loss Rafailov et al., 2023
SimPO simpo_loss Meng et al., 2024
ORPO orpo_loss Hong et al., 2024
GRPO grpo_loss Shao et al., 2024
PPO ppo_clip_loss Schulman et al., 2017
REINFORCE reinforce_loss Williams, 1992

Loss Dispatch

create_alignment_loss_fn(*, method, ref_model=None, beta=0.1, kl_coeff=0.04, lambda_weight=1.0, gamma=0.5, clip_eps=0.2)

Create a loss function for the given alignment method.

Returns a callable (model, batch, outputs) -> loss that handles forward passes on chosen/rejected pairs and reference model inference.

Source code in xaytune/recipes/align/loss_dispatch.py 
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def create_alignment_loss_fn(
    *,
    method: str,
    ref_model: Any | None = None,
    beta: float = 0.1,
    kl_coeff: float = 0.04,
    lambda_weight: float = 1.0,
    gamma: float = 0.5,
    clip_eps: float = 0.2,
) -> Callable[..., torch.Tensor]:
    """Create a loss function for the given alignment method.

    Returns a callable ``(model, batch, outputs) -> loss`` that handles
    forward passes on chosen/rejected pairs and reference model inference.
    """

    def loss_fn(
        model: Any,
        batch: dict[str, Any],
        outputs: Any,
    ) -> torch.Tensor:
        if not _has_alignment_fields(method, batch):
            loss: torch.Tensor = outputs.loss if hasattr(outputs, "loss") else outputs
            return loss

        if method == "dpo":
            return _dpo_step(model, batch, ref_model, beta=beta)
        elif method == "grpo":
            return _grpo_step(model, batch, ref_model, kl_coeff=kl_coeff)
        elif method == "orpo":
            return _orpo_step(model, batch, outputs, lambda_weight=lambda_weight)
        elif method == "simpo":
            return _simpo_step(model, batch, beta=beta, gamma=gamma)
        elif method == "ppo":
            return _ppo_step(model, batch, clip_eps=clip_eps)
        elif method == "reinforce":
            return _reinforce_step(model, batch)
        else:
            raise ValueError(f"Unknown alignment method: {method}")

    return loss_fn

is_alignment_method(method)

Return whether method is a known alignment method.

Source code in xaytune/recipes/align/loss_dispatch.py 
31
32
33
def is_alignment_method(method: str) -> bool:
    """Return whether *method* is a known alignment method."""
    return method in ALIGNMENT_METHODS

DPO

dpo_loss(*, policy_chosen_logps, policy_rejected_logps, ref_chosen_logps, ref_rejected_logps, beta=0.1)

Compute Direct Preference Optimization loss (Rafailov et al., 2023).

Source code in xaytune/recipes/align/dpo.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
def dpo_loss(
    *,
    policy_chosen_logps: torch.Tensor,
    policy_rejected_logps: torch.Tensor,
    ref_chosen_logps: torch.Tensor,
    ref_rejected_logps: torch.Tensor,
    beta: float = 0.1,
) -> torch.Tensor:
    """Compute Direct Preference Optimization loss (Rafailov et al., 2023)."""
    chosen_rewards = beta * (policy_chosen_logps - ref_chosen_logps)
    rejected_rewards = beta * (policy_rejected_logps - ref_rejected_logps)

    logits = chosen_rewards - rejected_rewards

    return -F.logsigmoid(logits).mean()

SimPO

simpo_loss(*, policy_chosen_logps, policy_rejected_logps, chosen_lengths, rejected_lengths, beta=2.0, gamma=0.5)

Compute Simple Preference Optimization loss (Meng et al., 2024).

Source code in xaytune/recipes/align/simpo.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
def simpo_loss(
    *,
    policy_chosen_logps: torch.Tensor,
    policy_rejected_logps: torch.Tensor,
    chosen_lengths: torch.Tensor,
    rejected_lengths: torch.Tensor,
    beta: float = 2.0,
    gamma: float = 0.5,
) -> torch.Tensor:
    """Compute Simple Preference Optimization loss (Meng et al., 2024)."""
    chosen_avg = policy_chosen_logps / chosen_lengths.float().clamp(min=1)
    rejected_avg = policy_rejected_logps / rejected_lengths.float().clamp(min=1)

    logits = beta * (chosen_avg - rejected_avg) - gamma

    return -F.logsigmoid(logits).mean()

ORPO

orpo_loss(*, sft_loss, policy_chosen_logps, policy_rejected_logps, lambda_weight=1.0)

Compute Odds Ratio Preference Optimization loss (Hong et al., 2024).

Source code in xaytune/recipes/align/orpo.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
def orpo_loss(
    *,
    sft_loss: torch.Tensor,
    policy_chosen_logps: torch.Tensor,
    policy_rejected_logps: torch.Tensor,
    lambda_weight: float = 1.0,
) -> torch.Tensor:
    """Compute Odds Ratio Preference Optimization loss (Hong et al., 2024)."""
    log_odds_ratio = (policy_chosen_logps - policy_rejected_logps) - (
        torch.log1p(-policy_chosen_logps.exp()) - torch.log1p(-policy_rejected_logps.exp())
    )

    or_loss = -F.logsigmoid(log_odds_ratio).mean()

    return sft_loss + lambda_weight * or_loss

GRPO

grpo_loss(*, logprobs, ref_logprobs=None, advantages, kl_coeff=0.04)

Compute Group Relative Policy Optimization loss (Shao et al., 2024).

Source code in xaytune/recipes/align/grpo.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def grpo_loss(
    *,
    logprobs: torch.Tensor,
    ref_logprobs: torch.Tensor | None = None,
    advantages: torch.Tensor,
    kl_coeff: float = 0.04,
) -> torch.Tensor:
    """Compute Group Relative Policy Optimization loss (Shao et al., 2024)."""
    policy_loss = -(logprobs * advantages).mean()

    if ref_logprobs is not None and kl_coeff > 0:
        kl = (logprobs - ref_logprobs).mean()
        return policy_loss + kl_coeff * kl

    return policy_loss

compute_group_advantages(rewards)

Normalize rewards to zero-mean unit-variance advantages.

Source code in xaytune/recipes/align/grpo.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
def compute_group_advantages(rewards: torch.Tensor) -> torch.Tensor:
    """Normalize rewards to zero-mean unit-variance advantages."""
    if rewards.numel() <= 1:
        return torch.zeros_like(rewards)

    mean = rewards.mean()
    std = rewards.std()

    if std < 1e-8:
        return torch.zeros_like(rewards)

    return (rewards - mean) / (std + 1e-8)

PPO / REINFORCE

ppo_clip_loss(*, logprobs, old_logprobs, advantages, clip_eps=0.2)

Compute the clipped surrogate policy gradient objective.

This implements only the clipped loss term from PPO (Schulman et al., 2017). It does NOT include rollout buffers, GAE, value model training, or multiple optimization epochs. See module docstring for details.

Source code in xaytune/recipes/align/ppo.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
def ppo_clip_loss(
    *,
    logprobs: torch.Tensor,
    old_logprobs: torch.Tensor,
    advantages: torch.Tensor,
    clip_eps: float = 0.2,
) -> torch.Tensor:
    """Compute the clipped surrogate policy gradient objective.

    This implements only the clipped loss term from PPO (Schulman et al., 2017).
    It does NOT include rollout buffers, GAE, value model training, or multiple
    optimization epochs. See module docstring for details.
    """
    ratio = torch.exp(logprobs - old_logprobs)

    unclipped = ratio * advantages
    clipped = torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps) * advantages

    return -torch.min(unclipped, clipped).mean()

ppo_value_loss(*, values, returns)

Compute value function MSE loss (used alongside the clipped policy gradient).

Source code in xaytune/recipes/align/ppo.py 
36
37
38
39
40
41
42
def ppo_value_loss(
    *,
    values: torch.Tensor,
    returns: torch.Tensor,
) -> torch.Tensor:
    """Compute value function MSE loss (used alongside the clipped policy gradient)."""
    return (values - returns).pow(2).mean()

reinforce_loss(*, logprobs, advantages)

Compute REINFORCE policy gradient loss.

Source code in xaytune/recipes/align/ppo.py 
45
46
47
48
49
50
51
def reinforce_loss(
    *,
    logprobs: torch.Tensor,
    advantages: torch.Tensor,
) -> torch.Tensor:
    """Compute REINFORCE policy gradient loss."""
    return -(logprobs * advantages).mean()

Log-Probabilities

get_per_token_logps(logits, labels)

Compute per-token log probabilities from logits and label ids.

Source code in xaytune/recipes/align/logprobs.py 
 7
 8
 9
10
11
12
13
14
def get_per_token_logps(
    logits: torch.Tensor,
    labels: torch.Tensor,
) -> torch.Tensor:
    """Compute per-token log probabilities from logits and label ids."""
    log_probs = F.log_softmax(logits[:, :-1, :], dim=-1)
    target = labels[:, 1:]
    return torch.gather(log_probs, dim=2, index=target.unsqueeze(2)).squeeze(2)

get_sequence_logps(logits, labels, mask=None, prompt_length=0)

Sum per-token log probabilities into a sequence-level log probability.

When prompt_length is provided, tokens before that position are excluded from the sum so only response tokens contribute.

Source code in xaytune/recipes/align/logprobs.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def get_sequence_logps(
    logits: torch.Tensor,
    labels: torch.Tensor,
    mask: torch.Tensor | None = None,
    prompt_length: torch.Tensor | int = 0,
) -> torch.Tensor:
    """Sum per-token log probabilities into a sequence-level log probability.

    When ``prompt_length`` is provided, tokens before that position are
    excluded from the sum so only response tokens contribute.
    """
    per_token = get_per_token_logps(logits, labels)
    if mask is not None:
        per_token = per_token * mask[:, 1:]
    if isinstance(prompt_length, int) and prompt_length > 0:
        per_token[:, :prompt_length] = 0.0
    elif isinstance(prompt_length, torch.Tensor) and prompt_length.any():
        for i, pl in enumerate(prompt_length):
            if pl > 0:
                per_token[i, : int(pl.item())] = 0.0
    return per_token.sum(dim=-1)

get_model_logps(model, input_ids, attention_mask=None, labels=None)

Run a forward pass and return sequence log probabilities (no grad).

Source code in xaytune/recipes/align/logprobs.py 
40
41
42
43
44
45
46
47
48
49
50
51
def get_model_logps(
    model: torch.nn.Module,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    labels: torch.Tensor | None = None,
) -> torch.Tensor:
    """Run a forward pass and return sequence log probabilities (no grad)."""
    if labels is None:
        labels = input_ids
    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    return get_sequence_logps(outputs.logits, labels, attention_mask)

Rewards

default_reward(prompt, response)

Baseline reward that always returns 0.

Source code in xaytune/recipes/align/rewards.py 
 98
 99
100
101
@register_reward("default")
def default_reward(prompt: str, response: str) -> float:
    """Baseline reward that always returns 0."""
    return 0.0

length_penalty_reward(prompt, response, *, target_length=200, penalty_scale=0.001)

Penalize responses that deviate from target_length characters.

Source code in xaytune/recipes/align/rewards.py 
104
105
106
107
108
109
110
111
112
113
114
@register_reward("length_penalty")
def length_penalty_reward(
    prompt: str,
    response: str,
    *,
    target_length: int = 200,
    penalty_scale: float = 0.001,
) -> float:
    """Penalize responses that deviate from *target_length* characters."""
    diff = abs(len(response) - target_length)
    return -penalty_scale * diff

format_check_reward(prompt, response, *, required_markers=None)

Reward based on the fraction of required_markers present in the response.

Source code in xaytune/recipes/align/rewards.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
@register_reward("format_check")
def format_check_reward(
    prompt: str,
    response: str,
    *,
    required_markers: list[str] | None = None,
) -> float:
    """Reward based on the fraction of *required_markers* present in the response."""
    if required_markers is None:
        required_markers = []
    if not required_markers:
        return 0.0
    matched = sum(1 for m in required_markers if m in response)
    return matched / len(required_markers)

composite_reward(prompt, response, *, reward_names=None, weights=None)

Weighted combination of multiple registered reward functions.

Source code in xaytune/recipes/align/rewards.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
@register_reward("composite")
def composite_reward(
    prompt: str,
    response: str,
    *,
    reward_names: list[str] | None = None,
    weights: list[float] | None = None,
) -> float:
    """Weighted combination of multiple registered reward functions."""
    if not reward_names:
        return 0.0
    if weights is None:
        weights = [1.0] * len(reward_names)
    total = 0.0
    for name, weight in zip(reward_names, weights):
        fn = reward_registry.get(name)
        total += weight * fn(prompt, response)
    return total

Agent Rewards

Reward functions for agent alignment with GRPO/PPO. Score agent responses based on tool usage quality, task completion, and efficiency. All rewards use <tool_call> tag parsing with pluggable custom parsers.

# In training config:
online_rl:
  reward_name: agent_composite
  reward_kwargs:
    expected_tools: ["search", "calculator"]
    success_markers: ["Done"]
    max_steps: 5

tool_use_quality_reward(prompt, response, *, expected_tools=None, required_args=None, parser=None)

Reward based on using the expected tools with required arguments.

Parameters:

Name Type Description Default
prompt str

The input prompt

 required
response str

The agent's response

 required
expected_tools list[str] | None

List of tool names that should be used

 None
required_args dict[str, list[str]] | None

Dict mapping tool names to lists of required argument names

 None
parser Callable | None

Optional custom parser

 None

Returns:

Type Description
float

Score from 0.0 to 1.0
Source code in xaytune/recipes/align/agent_rewards.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
@register_reward("tool_use_quality")
def tool_use_quality_reward(
    prompt: str,
    response: str,
    *,
    expected_tools: list[str] | None = None,
    required_args: dict[str, list[str]] | None = None,
    parser: Callable | None = None,
) -> float:
    """Reward based on using the expected tools with required arguments.

    Args:
        prompt: The input prompt
        response: The agent's response
        expected_tools: List of tool names that should be used
        required_args: Dict mapping tool names to lists of required argument names
        parser: Optional custom parser

    Returns:
        Score from 0.0 to 1.0
    """
    calls = parse_tool_calls(response, parser=parser)

    if not calls:
        return 1.0 if not expected_tools else 0.0

    if not expected_tools:
        return 1.0

    called_names = {call.name for call in calls}
    matched = sum(1 for t in expected_tools if t in called_names)
    score_parts = [float(matched)]
    total_parts = len(expected_tools)

    if required_args:
        for call in calls:
            if call.name in required_args:
                required = required_args[call.name]
                present = sum(1 for a in required if a in call.arguments)
                total_parts += len(required)
                score_parts.append(float(present))

    if total_parts == 0:
        return 1.0

    return sum(score_parts) / total_parts

task_completion_reward(prompt, response, *, success_markers=None, failure_markers=None, parser=None)

Reward based on task completion indicators.

Parameters:

Name Type Description Default
prompt str

The input prompt

 required
response str

The agent's response

 required
success_markers list[str] | None

Phrases indicating successful completion

 None
failure_markers list[str] | None

Phrases indicating failure

 None
parser Callable | None

Optional custom parser (not used here)

 None

Returns:

Type Description
float

Score from 0.0 to 1.0
Source code in xaytune/recipes/align/agent_rewards.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
@register_reward("task_completion")
def task_completion_reward(
    prompt: str,
    response: str,
    *,
    success_markers: list[str] | None = None,
    failure_markers: list[str] | None = None,
    parser: Callable | None = None,
) -> float:
    """Reward based on task completion indicators.

    Args:
        prompt: The input prompt
        response: The agent's response
        success_markers: Phrases indicating successful completion
        failure_markers: Phrases indicating failure
        parser: Optional custom parser (not used here)

    Returns:
        Score from 0.0 to 1.0
    """
    if failure_markers:
        for marker in failure_markers:
            if marker.lower() in response.lower():
                return 0.0

    if success_markers:
        matched = sum(1 for m in success_markers if m.lower() in response.lower())
        return matched / len(success_markers)

    last_result = response.rfind("</tool_result>")
    if last_result == -1:
        return 1.0 if response.strip() else 0.0

    after_tools = response[last_result + len("</tool_result>") :].strip()
    return 1.0 if after_tools else 0.0

efficiency_reward(prompt, response, *, max_steps=10, optimal_steps=None, parser=None)

Reward based on efficiency (fewer tool calls is better).

Parameters:

Name Type Description Default
prompt str

The input prompt

 required
response str

The agent's response

 required
max_steps int

Maximum acceptable number of tool calls

 10
optimal_steps int | None

Optimal number of tool calls (if known)

 None
parser Callable | None

Optional custom parser

 None

Returns:

Type Description
float

Score from 0.0 to 1.0
Source code in xaytune/recipes/align/agent_rewards.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
@register_reward("efficiency")
def efficiency_reward(
    prompt: str,
    response: str,
    *,
    max_steps: int = 10,
    optimal_steps: int | None = None,
    parser: Callable | None = None,
) -> float:
    """Reward based on efficiency (fewer tool calls is better).

    Args:
        prompt: The input prompt
        response: The agent's response
        max_steps: Maximum acceptable number of tool calls
        optimal_steps: Optimal number of tool calls (if known)
        parser: Optional custom parser

    Returns:
        Score from 0.0 to 1.0
    """
    calls = parse_tool_calls(response, parser=parser)
    num_calls = len(calls)

    if num_calls == 0:
        return 1.0 if response.strip() else 0.0

    if optimal_steps is not None:
        diff = abs(num_calls - optimal_steps)
        return max(0.0, 1.0 - diff / max_steps)

    return max(0.0, 1.0 - num_calls / max_steps)

agent_composite_reward(prompt, response, *, quality_weight=0.4, completion_weight=0.4, efficiency_weight=0.2, parser=None, expected_tools=None, required_args=None, success_markers=None, failure_markers=None, max_steps=10, optimal_steps=None)

Weighted combination of tool_use_quality, task_completion, and efficiency.

Source code in xaytune/recipes/align/agent_rewards.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
@register_reward("agent_composite")
def agent_composite_reward(
    prompt: str,
    response: str,
    *,
    quality_weight: float = 0.4,
    completion_weight: float = 0.4,
    efficiency_weight: float = 0.2,
    parser: Callable | None = None,
    expected_tools: list[str] | None = None,
    required_args: dict[str, list[str]] | None = None,
    success_markers: list[str] | None = None,
    failure_markers: list[str] | None = None,
    max_steps: int = 10,
    optimal_steps: int | None = None,
) -> float:
    """Weighted combination of tool_use_quality, task_completion, and efficiency."""
    quality = tool_use_quality_reward(
        prompt,
        response,
        expected_tools=expected_tools,
        required_args=required_args,
        parser=parser,
    )
    completion = task_completion_reward(
        prompt,
        response,
        success_markers=success_markers,
        failure_markers=failure_markers,
        parser=parser,
    )
    eff = efficiency_reward(
        prompt,
        response,
        max_steps=max_steps,
        optimal_steps=optimal_steps,
        parser=parser,
    )
    return quality_weight * quality + completion_weight * completion + efficiency_weight * eff

parse_tool_calls(text, parser=None)

Parse tool calls from text containing tags.

Parameters:

Name Type Description Default
text str

Text potentially containing tool calls

 required
parser Callable[[str], list[ParsedToolCall]] | None

Optional custom parser function

 None

Returns:

Type Description
list[ParsedToolCall]

List of ParsedToolCall objects
Source code in xaytune/recipes/align/agent_rewards.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def parse_tool_calls(
    text: str,
    parser: Callable[[str], list[ParsedToolCall]] | None = None,
) -> list[ParsedToolCall]:
    """Parse tool calls from text containing <tool_call> tags.

    Args:
        text: Text potentially containing tool calls
        parser: Optional custom parser function

    Returns:
        List of ParsedToolCall objects
    """
    if parser is not None:
        return parser(text)

    calls: list[ParsedToolCall] = []
    pattern = r"<tool_call>\s*(\{.*?\})\s*</tool_call>"

    for match in re.finditer(pattern, text, re.DOTALL):
        try:
            data = json.loads(match.group(1))
            name = data.get("name", "")
            arguments = data.get("arguments", {})
            calls.append(ParsedToolCall(name=name, arguments=arguments))
        except (json.JSONDecodeError, AttributeError):
            continue

    return calls

ParsedToolCall(name, arguments) dataclass

Represents a parsed tool call from agent output.