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

Code before:
          lambda leaf, indices=batch_indices: leaf[indices], rollout_batch
      )
      yield mini_batch

  def __iter__(self):
    """The main generator for the iterable."""
    for _ in range(self.repeat):
      for rollout_batch in self._data:
        yield from self._shuffle_and_slice_one_batch(rollout_batch)


@flax.struct.dataclass(frozen=True)
class TrainExample:
  prompt_ids: jax.Array
  prompt_mask: jax.Array
  completion_ids: jax.Array
  completion_mask: jax.Array
  advantages: jax.Array
  ref_per_token_logps: jax.Array | None
  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":

Review comment (suggestion):
Remove `None`, change to

```
method: str = "low_var_kl"
```

Code after:
          lambda leaf, indices=batch_indices: leaf[indices], rollout_batch
      )
      yield mini_batch

  def __iter__(self):
    """The main generator for the iterable."""
    for _ in range(self.repeat):
      for rollout_batch in self._data:
        yield from self._shuffle_and_slice_one_batch(rollout_batch)


@flax.struct.dataclass(frozen=True)
class TrainExample:
  prompt_ids: jax.Array
  prompt_mask: jax.Array
  completion_ids: jax.Array
  completion_mask: jax.Array
  advantages: jax.Array
  ref_per_token_logps: jax.Array | None
  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