[LLM] Wire MATH and Countdown into GRPO and Expert Iteration scripts

[vmoens]

Stack from ghstack (oldest at bottom):

• -> [LLM] Wire MATH and Countdown into GRPO and Expert Iteration scripts #3546
• [LLM] Add Countdown numbers-game environment #3545
• [LLM] Add MATH (competition mathematics) environment #3544
• [LLM] Simplify IFEval reward aggregator #3543
• [LLM] Rewrite GSM8K reward function to follow standard GRPO conventions #3542

Add GRPO config files for math and countdown datasets, and update
grpo_utils.py and ei_utils.py to support the new dataset choices.

All four datasets (gsm8k, ifeval, math, countdown) are now selectable
via the env.dataset config key.

Made-with: Cursor

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3546

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 4 New Failures, 1 Cancelled Job, 6 Unrelated Failures

As of commit 783db8a with merge base 491ed0f ():

NEW FAILURES - The following jobs have failed:

• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  The version '3.10' with architecture 'x64' was not found for this operating system.
• Generate documentation / build-docs (3.12, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 2d1cf5d922c1d56bf430662e88109d16bb0dbd06add86122dcba6d63007003a8 /exec failed with exit code 2
• PR Label / add-label (gh)
  Process completed with exit code 1.
• Unit-tests on Linux / tests-cpu (3.14) / linux-job (gh)
  RuntimeError: Command docker exec -t 63f03eaa891db6e3f2e4ab4eeeca1f5fd2e07653cfca1eda046813709fb1a2db /exec failed with exit code 245

CANCELLED JOB - The following job was cancelled. Please retry:

• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
  ##[error]The operation was canceled.

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Unit-tests on Linux / tests-cpu (3.10) / linux-job (gh) (trunk failure)
  test/test_rb.py::TestStorages::test__rand_given_ndim_recompile
• Unit-tests on Linux / tests-cpu (3.11) / linux-job (gh) (trunk failure)
  test/test_rb.py::TestStorages::test__rand_given_ndim_recompile
• Unit-tests on Linux / tests-cpu (3.12) / linux-job (gh) (trunk failure)
  test/test_rb.py::TestStorages::test__rand_given_ndim_recompile
• Unit-tests on Linux / tests-cpu (3.13) / linux-job (gh) (trunk failure)
  test/test_rb.py::TestStorages::test__rand_given_ndim_recompile
• Unit-tests on Linux / tests-olddeps (3.10, 11.8) / linux-job (gh) (trunk failure)
  test/test_loggers.py::TestCSVLogger::test_log_video[mp4-steps1]
• Unit-tests on Linux / tests-optdeps (3.12, 13.0) / linux-job (gh) (trunk failure)
  test/test_rb.py::TestStorages::test__rand_given_ndim_recompile

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

⚠️ PR Title Label Error

Unknown or invalid prefix [LLM].

Current title: [LLM] Wire MATH and Countdown into GRPO and Expert Iteration scripts

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix	Label Applied	Example
[BugFix]	BugFix	[BugFix] Fix memory leak in collector
[Feature]	Feature	[Feature] Add new optimizer
[Doc] or [Docs]	Documentation	[Doc] Update installation guide
[Refactor]	Refactoring	[Refactor] Clean up module imports
[CI]	CI	[CI] Fix workflow permissions
[Test] or [Tests]	Tests	[Tests] Add unit tests for buffer
[Environment] or [Environments]	Environments	[Environments] Add Gymnasium support
[Data]	Data	[Data] Fix replay buffer sampling
[Performance] or [Perf]	Performance	[Performance] Optimize tensor ops
[BC-Breaking]	bc breaking	[BC-Breaking] Remove deprecated API
[Deprecation]	Deprecation	[Deprecation] Mark old function
[Quality]	Quality	[Quality] Fix typos and add codespell

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

[github-actions]

⚠️ PR Title Label Error

Unknown or invalid prefix [LLM].

Current title: [LLM] Wire MATH and Countdown into GRPO and Expert Iteration scripts

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix	Label Applied	Example
[BugFix]	BugFix	[BugFix] Fix memory leak in collector
[Feature]	Feature	[Feature] Add new optimizer
[Doc] or [Docs]	Documentation	[Doc] Update installation guide
[Refactor]	Refactoring	[Refactor] Clean up module imports
[CI]	CI	[CI] Fix workflow permissions
[Test] or [Tests]	Tests	[Tests] Add unit tests for buffer
[Environment] or [Environments]	Environments	[Environments] Add Gymnasium support
[Data]	Data	[Data] Fix replay buffer sampling
[Performance] or [Perf]	Performance	[Performance] Optimize tensor ops
[BC-Breaking]	bc breaking	[BC-Breaking] Remove deprecated API
[Deprecation]	Deprecation	[Deprecation] Mark old function
[Quality]	Quality	[Quality] Fix typos and add codespell

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).