trinity.algorithm.kl_fn package

Submodules

• trinity.algorithm.kl_fn.kl_fn module
  • KLFn
    • KLFn.__init__()
    • KLFn.update_kl_coef()
    • KLFn.apply_kl_penalty_to_reward()
    • KLFn.calculate_kl_loss()
    • KLFn.calculate_kl()
    • KLFn.default_args()
  • DummyKLFn
    • DummyKLFn.calculate_kl()
    • DummyKLFn.apply_kl_penalty_to_reward()
    • DummyKLFn.calculate_kl_loss()
  • K1Fn
    • K1Fn.calculate_kl()
  • K2Fn
    • K2Fn.calculate_kl()
  • K3Fn
    • K3Fn.calculate_kl()
  • AbsFn
    • AbsFn.calculate_kl()

Module contents

class trinity.algorithm.kl_fn.KLFn(adaptive: bool = False, kl_coef: float = 0.001, target_kl: float | None = None, horizon: float | None = None)

Bases: ABC

KL penalty and loss.

__init__(adaptive: bool = False, kl_coef: float = 0.001, target_kl: float | None = None, horizon: float | None = None) → None

update_kl_coef(current_kl: float, batch_size: int) → None

Update kl coefficient.

apply_kl_penalty_to_reward(experiences: Any) → Tuple[Any, Dict]

Apply KL penalty to reward. Only support DataProto input for now.

calculate_kl_loss(logprob: Tensor, ref_logprob: Tensor, response_mask: Tensor) → Tuple[Tensor, Dict]

Compute KL loss.

abstract calculate_kl(logprob: Tensor, ref_logprob: Tensor) → Tensor

Compute KL divergence between logprob and ref_logprob.

classmethod default_args()

Get the default initialization arguments.