Repository: google/tunix
File: tunix/rl/common.py (Python)

Code before:
  old_per_token_logps: jax.Array | None


def compute_kl_divergence(
    per_token_logps: jax.Array, 
    ref_per_token_logps: jax.Array, 
    method: str | None = None
) -> jax.Array:
  """Compute per token KL divergence between trained and reference policy.
  kl: Unbiased, high-variance estimator - simple forward KL: logp - ref_logp
  mse_kl: Biased, low-variance estimator - squared log-difference: 0.5 * (logp - ref_logp)^2
  low_var_kl: Unbiased, low-variance estimator - J. Schulman low-variance approx: (r - 1) - log r, where r = q/p = exp(ref_logp - logp)

  references:
  J. Schulman KL approx: http://joschu.net/blog/kl-approx.html
  Verl code: https://github.com/volcengine/verl/blob/f9035b70166b102de5e84c1819e2172b64a0ae30/verl/trainer/ppo/core_algos.py#L1323-L1358

  Args:
    per_token_logps: Per token log probabilities from the trained policy.
    ref_per_token_logps: Per token log probabilities from the reference policy.
    method: KL penalty method. If None, defaults to "low_var_kl"..

  Returns:
    KL divergence.
  """
  delta = per_token_logps - ref_per_token_logps

  if method == "kl":
    # Unbiased but high variance.
    return delta

  if method == "mse_kl":
    # Biased but low variance.
    return 0.5 * jnp.square(delta)

  if method == "low_var_kl":
    # Default to low-variance, unbiased estimator.
    # (r - 1) - log(r) = exp(ref_logp - logp) - (ref_logp - logp) - 1
    return jnp.exp(ref_per_token_logps - per_token_logps) \
           - (ref_per_token_logps - per_token_logps) - 1
  
  # Default to low-variance, unbiased estimator.
  return jnp.exp(ref_per_token_logps - per_token_logps) \
           - (ref_per_token_logps - per_token_logps) - 1



def selective_log_softmax(logits: jax.Array, input_ids: jax.Array) -> jax.Array:
  """Compute the log probablity based on the input ids.

  Args:

Review comment (suggestion):
Let's add a check up here...

```
if method not in (...):
  raise ValueError("`method` must be one of ... Received: ...")
```

Code after:
  old_per_token_logps: jax.Array | None


def compute_kl_divergence(
    per_token_logps: jax.Array,
    ref_per_token_logps: jax.Array,
    method: str = "low_var_kl",
) -> jax.Array:
  """Compute per token KL divergence between trained and reference policy.

  Based on `method`, we compute one of three kinds of KL divergence:
  - "kl": Unbiased, high-variance estimator. Simple Forward KL:
    `logp - ref_logp`.
  - "mse_kl": Biased, low-variance estimator. Squared log-difference:
    `0.5 * (logp - ref_logp)^2`.
  - "low_var_kl": Unbiased, low-variance estimator. J. Schulman low-variance
    approx: `(r - 1) - log r`, where `r = q/p = exp(ref_logp - logp)`.

  Args:
    per_token_logps: Per token log probabilities from the trained policy.
    ref_per_token_logps: Per token log probabilities from the reference policy.
    method: KL penalty method. Defaults to "low_var_kl".

  Returns:
    KL divergence.
  """
  if method == "kl":
    return per_token_logps - ref_per_token_logps
  elif method == "mse_kl":
    return 0.5 * jnp.square(per_token_logps - ref_per_token_logps)
  elif method == "low_var_kl":
    kl = ref_per_token_logps - per_token_logps
    return jnp.exp(kl) - (kl) - 1
  else:
    raise ValueError(
        "`method` must be one of 'kl', 'mse_kl', 'low_var_kl'. Received:"
        f" {method}"
    )


def selective_log_softmax(logits: jax.Array, input_ids: jax.Array) -> jax.Array:
  """Compute the log probablity based on the input ids.

  Args:
    logits: Logits from the model.
    input_ids: Input ids to get logits.

  Returns:
    Selected log probabilities.
  """
  logps = jax.nn.log_softmax(logits, axis=-1)