feat: add AIME-2026 benchmark.

docs/guides/eval.md

@@ -103,7 +103,7 @@ score=0.1000 (3.0/30)
 
 ## List of currently supported benchmarks
 
-- [AIME-2024 and AIME-2025](../../nemo_rl/data/datasets/eval_datasets/aime.py): the corresponding `data.dataset_name` are `"aime2024"` and `"aime2025"`.
+- [AIME-2024, AIME-2025, and AIME-2026](../../nemo_rl/data/datasets/eval_datasets/aime.py): the corresponding `data.dataset_name` are `"aime2024"`, `"aime2025"`, and `"aime2026"`.
 - [GPQA and GPQA-diamond](../../nemo_rl/data/datasets/eval_datasets/gpqa.py): the corresponding `data.dataset_name` are `"gpqa"` and `"gpqa_diamond"`.
 - [MATH and MATH-500](../../nemo_rl/data/datasets/eval_datasets/math.py): the corresponding `data.dataset_name` are `"math"` and `"math500"`.
 - [MMLU](../../nemo_rl/data/datasets/eval_datasets/mmlu.py): this also includes MMMLU (Multilingual MMLU), a total of 14 languages. When `data.dataset_name` is set to `mmlu`, the English version is used. If one wants to run evaluation on another language, `data.dataset_name` should be set to `mmlu_{language}` where `language` is one of following 14 values, `["AR-XY", "BN-BD", "DE-DE", "ES-LA", "FR-FR", "HI-IN", "ID-ID", "IT-IT", "JA-JP", "KO-KR", "PT-BR", "ZH-CN", "SW-KE", "YO-NG"]`.

nemo_rl/data/datasets/eval_datasets/__init__.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from nemo_rl.data.datasets.eval_datasets.aime import AIMEDataset
+from typing import cast
+
+from nemo_rl.data.datasets.eval_datasets.aime import AIMEDataset, AIMEVariant
 from nemo_rl.data.datasets.eval_datasets.gpqa import GPQADataset
 from nemo_rl.data.datasets.eval_datasets.local_math_dataset import LocalMathDataset
 from nemo_rl.data.datasets.eval_datasets.math import MathDataset
@@ -53,15 +55,9 @@ def load_eval_dataset(data_config):
             system_prompt_file=data_config["system_prompt_file"],
         )
     # aime
-    elif dataset_name == "aime2024":
-        base_dataset = AIMEDataset(
-            variant="2024",
-            prompt_file=data_config["prompt_file"],
-            system_prompt_file=data_config["system_prompt_file"],
-        )
-    elif dataset_name == "aime2025":
+    elif dataset_name in ["aime2024", "aime2025", "aime2026"]:
         base_dataset = AIMEDataset(
-            variant="2025",
+            variant=cast(AIMEVariant, dataset_name[4:]),
             prompt_file=data_config["prompt_file"],
             system_prompt_file=data_config["system_prompt_file"],
         )

nemo_rl/data/datasets/eval_datasets/aime.py

@@ -21,11 +21,13 @@
 from nemo_rl.data import processors
 from nemo_rl.data.interfaces import TaskDataSpec
 
+AIMEVariant = Literal["2024", "2025", "2026"]
+
 
 class AIMEDataset:
     def __init__(
         self,
-        variant: Literal["2024", "2025"] = "2025",
+        variant: AIMEVariant = "2025",
         prompt_file: Optional[str] = None,
         system_prompt_file: Optional[str] = None,
     ):
@@ -37,6 +39,9 @@ def __init__(
             ds1 = load_dataset("opencompass/AIME2025", "AIME2025-II", split="test")
             ds = concatenate_datasets([ds0, ds1])
             self.input_key = "question"
+        elif variant == "2026":
+            ds = load_dataset("MathArena/aime_2026", split="train")
+            self.input_key = "problem"
         else:
             raise ValueError(f"Invalid variant for aime dataset: aime{variant}")
 

tests/unit/data/datasets/test_eval_dataset.py

@@ -82,6 +82,39 @@ def test_math_dataset():
         )
 
 
+@pytest.mark.parametrize("variant", ["aime2024", "aime2025", "aime2026"])
+@pytest.mark.skip(reason="dataset download is flaky")
+def test_aime_dataset(variant):
+    # load the dataset
+    data_config = {
+        "dataset_name": variant,
+        "prompt_file": None,
+        "system_prompt_file": None,
+    }
+    aime_dataset = load_eval_dataset(data_config)
+
+    # load the tokenizer
+    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct")
+
+    # check that the dataset is formatted correctly
+    for example in aime_dataset.rekeyed_ds.take(5):
+        assert "problem" in example
+        assert "expected_answer" in example
+
+        ## check that applying chat template works as expected
+        default_templated = tokenizer.apply_chat_template(
+            [{"role": "user", "content": example["problem"]}],
+            tokenize=False,
+            add_generation_prompt=False,
+            add_special_tokens=False,
+        )
+
+        assert (
+            default_templated
+            == f"<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n{example['problem']}<|im_end|>\n"
+        )
+
+
 @pytest.mark.skip(reason="dataset download is flaky")
 def test_mmlu_dataset():
     # load the dataset