feat: additional kl metrics

[ZhiyuLi-Nvidia]

What does this PR do ?

Add metrics:

• https://wandb.ai/nvidia/grpo-dev-zhiyul/runs/kgjjunme?nw=poxtgkh2ggc
• https://wandb.ai/nvidia/zhiyul-nano-v3/runs/yeulq88w?nw=pty704q4dt

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

Summary by CodeRabbit

Release Notes

• New Features

  • Added three divergence metrics (gen_kl, policy_kl, and js_divergence) to track differences between model generation and training distributions during reinforcement learning training.

• Documentation

  • Added comprehensive guide explaining KL Divergence metrics, including definitions, interpretation guidance, and practical examples for monitoring model behavior.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR adds KL divergence metrics (gen_kl, policy_kl, and js_divergence) to track policy distribution divergence. These metrics are implemented in the ClippedPGLossFn loss function and documented with detailed explanations, interpretations, and concrete examples in the guides.

Changes

Cohort / File(s)	Summary
Documentation Enhancement docs/guides/grpo.md	Added comprehensive KL Divergence subsection under Multiplicative Token Probability Error metrics, defining gen_kl, policy_kl, and js_divergence with reference distributions, concrete examples, expected ranges, and guidance on interpretation and spike investigation.
Metrics Implementation nemo_rl/algorithms/loss_functions.py	Added three new divergence metrics to ClippedPGLossFn.call: gen_kl (forward KL), policy_kl (reverse KL), and js_divergence (Jensen–Shannon), computed from prev_logprobs and generation_logprobs with masked normalization, and integrated into the returned metrics dictionary.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Pre-merge checks and finishing touches

✅ Passed checks (4 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Test Results For Major Changes	✅ Passed	This PR adds three new observational metrics (gen_kl, policy_kl, and js_divergence) to the ClippedPGLossFn loss function's returned metrics dictionary and updates documentation with a KL divergence subsection. According to the provided analysis, these metrics are purely monitoring additions that do not modify core loss computation or training algorithm logic. The PR description lacks test results documentation and shows all pre-checks unchecked. However, since these metrics do not affect numerics, convergence, or performance—they only add new values to a monitoring dictionary—they constitute a minor enhancement rather than a major change that would require regression testing documentation.
Title Check	✅ Passed	The pull request title "feat: additional kl metrics" is directly and fully related to the main changes in the changeset. The PR introduces three new metrics—gen_kl, policy_kl, and js_divergence—to the ClippedPGLossFn loss function, along with documentation explaining these divergence-based metrics. The title is concise, clear, and specific enough to convey the primary change at a glance; it accurately captures that the PR adds new KL-related metrics to the system.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch zhiyul/kl_metrics

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between e762237 and fa2bff4.

📒 Files selected for processing (2)

• docs/guides/grpo.md (1 hunks)
• nemo_rl/algorithms/loss_functions.py (2 hunks)

🧰 Additional context used

📓 Path-based instructions (3)

docs/**/*.md

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

When a markdown doc under docs/**/*.md is added or renamed, update docs/index.md to include it in the appropriate section

Files:

• docs/guides/grpo.md

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Follow the Google Python Style Guide for all Python code
Target Python 3.12+ for all Python code in NeMo-RL
Indent Python code with 4 spaces; do not use tabs
Python filenames should be snake_case (e.g., some_file.py)
Class names should be PascalCase
Function and method names should be snake_case
Local variable names should be snake_case; if starting with a number, prefix with k (e.g., k_99th_percentile)
Global variables should be UPPER_SNAKE_CASE and prefixed with G_ (e.g., G_MY_GLOBAL)
Constants should be UPPER_SNAKE_CASE
Avoid shadowing variables declared in an outer scope
Initialize all externally visible members of a class in the constructor
For public interfaces used outside a file, prefer docstrings over comments
Use comments mainly for code within a function or interfaces local to a file
Commented-out code must include a nearby comment explaining usage and why it is commented out; otherwise remove before merging
Use Google-style docstrings for classes and functions (Sphinx-parseable)
Avoid using reflection when functionality can be easily achieved without it
Limit except clauses to the smallest specific set of exceptions possible
For duck-typing via try/except, keep the try body minimal and use else for main logic
Add the NVIDIA copyright header (with current year) at the top of all Python files, excluding tests/ and test-only scripts

Files:

• nemo_rl/algorithms/loss_functions.py

nemo_rl/**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

nemo_rl/**/*.py: Do not set non-None configuration defaults in code; YAML is the single source of truth for defaults
Access required config attributes directly (e.g., policy_cfg["precision"]) and assume presence; do not introduce hidden defaults
Express configuration optionality via TypedDict using typing.NotRequired
When adding a new config key to a TypedDict subclass, document the key’s purpose, valid values/types, and recommended default in code
For any class or function decorated with @ray.remote, add '# pragma: no cover' on the class/def line (and on remote functions)

Files:

• nemo_rl/algorithms/loss_functions.py

🧬 Code graph analysis (1)

nemo_rl/algorithms/loss_functions.py (1)

nemo_rl/algorithms/utils.py (1)

• masked_mean (134-146)

🪛 LanguageTool

docs/guides/grpo.md

[grammar] ~285-~285: Ensure spelling is correct
Context: ...: exp(-5 + 15) - (-5 + 15) - 1 = 22015, servere mismatch dominiting the metrics. * `js_...

(QB_NEW_EN_ORTHOGRAPHY_ERROR_IDS_1)

---

[grammar] ~285-~285: Ensure spelling is correct
Context: ...(-5 + 15) - 1 = 22015, servere mismatch dominiting the metrics. * js_divergence: close t...

(QB_NEW_EN_ORTHOGRAPHY_ERROR_IDS_1)

🪛 markdownlint-cli2 (0.18.1)

docs/guides/grpo.md

276-276: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

🪛 Ruff (0.14.1)

nemo_rl/algorithms/loss_functions.py

191-191: Comment contains ambiguous – (EN DASH). Did you mean - (HYPHEN-MINUS)?

(RUF003)

🔇 Additional comments (3)

nemo_rl/algorithms/loss_functions.py (3)

171-179: LGTM! Correct implementation of gen_kl.

The gen_kl metric correctly computes KL(P_gen || P_train) using the Schulman approximation. The formula exp(log_ratio) - log_ratio - 1 with log_ratio = prev_logprobs - generation_logprobs properly implements the divergence from the generation distribution to the training distribution.

---

181-189: LGTM! Correct implementation of policy_kl.

The policy_kl metric correctly computes KL(P_train || P_gen) by reversing the log ratio direction. This provides a complementary view to gen_kl, treating the policy distribution as ground truth.

---

416-418: LGTM! Metrics properly integrated into return dictionary.

The three new KL/divergence metrics are correctly added to the metrics dictionary and will be logged alongside existing metrics.

[terrykong]

small comment. otherwise lgtm