# 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.
import torch
import torch.nn.functional as F
[docs]
def compute_logprobs_from_logits(logits, target):
logprobs = -F.cross_entropy(logits, target=target, reduction="none") # [B, action-dim]
return logprobs
[docs]
def compute_entropy_from_logits(logits, epsilon=1e-10):
"""
Compute entropy by logits.
Args:
logits: [B, vocab-size, seq-len]
epsilon: Small constant for numerical stability.
Returns:
entropy: [B, seq-len]
"""
all_probs = F.softmax(logits, dim=1) # [B, vocab-size, seq-len]
all_log_probs = torch.log(all_probs + epsilon)
entropy = -torch.sum(all_probs * all_log_probs, dim=1) # [B, seq-len]
return entropy