# Copyright 2025 The RLinf Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Callable, Dict, Optional, Tuple
import torch
from rlightning.policy.utils.utils import huber_loss, masked_mean, masked_mean_ratio
[docs]
def compute_ppo_actor_critic_loss(**kwargs) -> Tuple[torch.Tensor, Dict]:
"""Compute PPO actor-critic loss.
Args:
**kwargs: Keyword arguments forwarded to actor and critic loss functions.
Returns:
Tuple[torch.Tensor, Dict]: Total loss and metrics dictionary.
"""
metrics_data = {}
actor_loss, actor_metrics_data = compute_ppo_actor_loss(**kwargs)
critic_loss, critic_metrics_data = compute_ppo_critic_loss(**kwargs)
loss = actor_loss + critic_loss
metrics_data.update(actor_metrics_data)
metrics_data.update(critic_metrics_data)
return loss, metrics_data
[docs]
def compute_ppo_actor_loss(
logprobs: torch.Tensor,
old_logprobs: torch.Tensor,
clip_ratio_low: float,
clip_ratio_high: float,
advantages: torch.Tensor,
loss_mask: Optional[torch.Tensor] = None,
c_clip: Optional[float] = None,
loss_agg_func: Optional[Callable[..., torch.Tensor]] = masked_mean,
max_episode_steps: Optional[int] = None,
loss_mask_sum: Optional[torch.Tensor] = None,
critic_warmup: Optional[bool] = False,
**kwargs,
) -> Tuple[torch.Tensor, Dict]:
"""Compute PPO actor loss.
Args:
logprobs (torch.FloatTensor): Log probabilities of actions.
old_logprobs (torch.FloatTensor): Old log probabilities of actions.
clip_ratio_low (float): Lower bound of clipping ratio.
clip_ratio_high (float): Upper bound of clipping ratio.
advantages (torch.FloatTensor): GAE (normalized) advantages.
loss_mask (Optional[torch.BoolTensor]): Mask for valid entries.
c_clip (Optional[float]): Optional clipping coefficient.
loss_agg_func (Optional[Callable[..., torch.Tensor]]): Aggregation function.
max_episode_steps (Optional[int]): Max episode length for normalization.
loss_mask_sum (Optional[torch.Tensor]): Sum of mask values for normalization.
critic_warmup (Optional[bool]): If True, zero out actor loss.
**kwargs: Unused extra keyword arguments.
Returns:
Tuple[torch.Tensor, Dict]: Actor loss and metrics dictionary.
"""
loss_mask_ratio = None
if max_episode_steps is not None and loss_mask_sum is not None and loss_mask is not None:
loss_mask_ratio = (loss_mask_sum * 1.0) / max_episode_steps
loss_agg_func = masked_mean_ratio
if loss_mask is None:
loss_mask = torch.ones_like(logprobs).bool()
assert logprobs.dtype == torch.float32
assert old_logprobs.dtype == torch.float32
assert advantages.dtype == torch.float32
loss_mask_count = loss_mask.count_nonzero() or 1
# For numerical stability.
ratio = torch.where(loss_mask, torch.exp(logprobs - old_logprobs), 0)
approx_kl = torch.where(loss_mask, (logprobs - old_logprobs).detach(), 0.0)
clipped_ratio = torch.clamp(ratio, 1.0 - clip_ratio_low, 1.0 + clip_ratio_high)
policy_loss1 = -advantages * ratio
policy_loss2 = -advantages * clipped_ratio
clip_mask = policy_loss1.detach() < policy_loss2.detach()
policy_loss = torch.max(policy_loss1, policy_loss2)
if c_clip is not None:
assert c_clip > 1.0, c_clip
policy_loss3 = torch.sign(advantages) * c_clip * advantages
dual_clip_mask = policy_loss3.detach() < policy_loss.detach()
policy_loss = torch.min(policy_loss, policy_loss3)
else:
dual_clip_mask = torch.zeros_like(clip_mask)
policy_loss = loss_agg_func(
policy_loss, loss_mask, loss_mask_ratio
) # default max_episode_steps is None
clip_mask = policy_loss1.detach() < policy_loss2.detach()
dual_clip_mask.logical_and_(loss_mask)
clip_fraction = clip_mask.logical_and_(loss_mask).count_nonzero() / loss_mask_count
approx_kl = -approx_kl.sum() / loss_mask_count
dual_cliped_ratio = torch.where(dual_clip_mask, ratio, 0)
if critic_warmup:
policy_loss = torch.tensor(0.0, device=policy_loss.device)
# Compile metrics for logging
metrics_data = {
"actor/policy_loss": policy_loss.detach(),
"actor/ratio": masked_mean(ratio.detach(), loss_mask),
"actor/clipped_ratio": masked_mean(clipped_ratio.detach(), loss_mask),
"actor/dual_cliped_ratio": masked_mean(dual_cliped_ratio.detach(), loss_mask),
"actor/approx_kl": approx_kl.detach(),
"actor/clip_fraction": clip_fraction.detach(),
}
return policy_loss, metrics_data
[docs]
def compute_ppo_critic_loss(
values: torch.Tensor,
returns: torch.Tensor,
prev_values: torch.Tensor,
value_clip: float,
huber_delta: float,
loss_mask: Optional[torch.Tensor] = None,
max_episode_steps: Optional[int] = None,
loss_mask_sum: Optional[torch.Tensor] = None,
**kwargs,
) -> Tuple[torch.Tensor, Dict]:
"""Compute PPO critic loss.
Args:
values (torch.Tensor): Current value predictions.
returns (torch.Tensor): Return values.
prev_values (torch.Tensor): Previous value predictions.
value_clip (float): Value clipping threshold.
huber_delta (float): Huber loss delta parameter.
loss_mask (Optional[torch.Tensor]): Mask for valid entries.
max_episode_steps (Optional[int]): Max episode length for normalization.
loss_mask_sum (Optional[torch.Tensor]): Sum of mask values for normalization.
**kwargs: Unused extra keyword arguments.
Returns:
Tuple[torch.Tensor, Dict]: Critic loss and metrics dictionary.
"""
loss_mask_ratio = None
loss_agg_func = masked_mean
if max_episode_steps is not None and loss_mask_sum is not None and loss_mask is not None:
loss_mask_ratio = (loss_mask_sum * 1.0) / max_episode_steps
loss_agg_func = masked_mean_ratio
value_pred_clipped = prev_values + (values - prev_values).clamp(
-value_clip, value_clip
) # [bsz, ] | [bsz, chunk-step]
value_loss_original = huber_loss(returns - values, huber_delta) # [bsz, ] | [bsz, chunk-step]
value_loss_clipped = huber_loss(
returns - value_pred_clipped, huber_delta
) # [bsz, ] | [bsz, chunk-step]
value_loss = torch.max(value_loss_original, value_loss_clipped)
value_loss = loss_agg_func(value_loss, loss_mask, loss_mask_ratio)
value_clip_indicator = (value_pred_clipped - prev_values).abs() > value_clip
value_clip_ratio = value_clip_indicator.float().mean()
# Compile metrics for logging
metrics_data = {
"critic/value_loss": value_loss.detach().item(),
"critic/value_clip_ratio": value_clip_ratio.detach().item(),
}
return value_loss, metrics_data