From b761dd14bfc8253f916467cdf897558d4ac52adc Mon Sep 17 00:00:00 2001
From: Terry Kong <terryk@nvidia.com>
Date: Mon, 2 Feb 2026 20:45:20 +0000
Subject: [PATCH] feat: add val_at_end for all algorithms

Signed-off-by: Terry Kong <terryk@nvidia.com>
---
 examples/configs/distillation_math.yaml       |   1 +
 examples/configs/dpo.yaml                     |   1 +
 examples/configs/grpo_math_1B.yaml            |   1 +
 examples/configs/grpo_math_1B_megatron.yaml   |   1 +
 examples/configs/rm.yaml                      |   1 +
 examples/configs/sft.yaml                     |   1 +
 examples/configs/sft_openmathinstruct2.yaml   |   1 +
 .../sft_openmathinstruct2_megatron.yaml       |   1 +
 examples/configs/vlm_grpo_3B.yaml             |   1 +
 examples/configs/vlm_grpo_3B_megatron.yaml    |   1 +
 ...rkplace_assistant_nemotron_nano_v2_9b.yaml |   1 +
 nemo_rl/algorithms/distillation.py            |  16 ++-
 nemo_rl/algorithms/dpo.py                     |  10 +-
 nemo_rl/algorithms/grpo.py                    |  22 +++-
 nemo_rl/algorithms/rm.py                      |  10 +-
 nemo_rl/algorithms/sft.py                     |  10 +-
 nemo_rl/utils/checkpoint.py                   |  30 ++++-
 .../configs/grpo_math_1B.yaml                 |   1 +
 tests/functional/test_converter_roundtrip.py  |   1 +
 tests/unit/algorithms/test_distillation.py    |   4 +
 tests/unit/algorithms/test_dpo.py             |   1 +
 tests/unit/algorithms/test_grpo.py            |   4 +
 tests/unit/algorithms/test_rm.py              |   1 +
 tests/unit/algorithms/test_sft.py             |   1 +
 tests/unit/utils/test_checkpoint.py           | 120 ++++++++++++++++++
 25 files changed, 222 insertions(+), 20 deletions(-)

diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml
index 078b2bef967..67ff8a71d2d 100644
--- a/examples/configs/distillation_math.yaml
+++ b/examples/configs/distillation_math.yaml
@@ -8,6 +8,7 @@ distillation:
     val_batch_size: 64
     val_period: 20
     val_at_start: false
+    val_at_end: false
     max_val_samples: 512
     topk_logits_k: 64
     seed: 42
diff --git a/examples/configs/dpo.yaml b/examples/configs/dpo.yaml
index f4b0c148161..0db99b8872f 100755
--- a/examples/configs/dpo.yaml
+++ b/examples/configs/dpo.yaml
@@ -7,6 +7,7 @@ dpo:
   val_global_batch_size: 8
   val_micro_batch_size: 1
   val_at_start: true
+  val_at_end: false
   seed: 42
 
   reference_policy_kl_penalty: 0.05
diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml
index 99bf8615f1f..0950de0cd53 100644
--- a/examples/configs/grpo_math_1B.yaml
+++ b/examples/configs/grpo_math_1B.yaml
@@ -9,6 +9,7 @@ grpo:
   use_leave_one_out_baseline: true
   val_period: 10
   val_at_start: false
+  val_at_end: false
   overlong_filtering: false
   max_val_samples: 256
   val_batch_size: 256
diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml
index bb51d509424..dfd02ecc779 100644
--- a/examples/configs/grpo_math_1B_megatron.yaml
+++ b/examples/configs/grpo_math_1B_megatron.yaml
@@ -10,6 +10,7 @@ grpo:
   use_leave_one_out_baseline: true
   val_period: 10
   val_at_start: false
+  val_at_end: false
   max_val_samples: 256
   val_batch_size: 256
   async_grpo:
diff --git a/examples/configs/rm.yaml b/examples/configs/rm.yaml
index 11e10167584..3051f4bc8f1 100644
--- a/examples/configs/rm.yaml
+++ b/examples/configs/rm.yaml
@@ -10,6 +10,7 @@ rm:
   val_global_batch_size: 32
   val_micro_batch_size: 1
   val_at_start: false
+  val_at_end: false
   seed: 42
 
 checkpointing:
diff --git a/examples/configs/sft.yaml b/examples/configs/sft.yaml
index 728b1c47ff8..2160abbb056 100644
--- a/examples/configs/sft.yaml
+++ b/examples/configs/sft.yaml
@@ -10,6 +10,7 @@ sft:
   val_global_batch_size: 32
   val_micro_batch_size: 1
   val_at_start: true
+  val_at_end: false
   seed: 42
 
 checkpointing:
diff --git a/examples/configs/sft_openmathinstruct2.yaml b/examples/configs/sft_openmathinstruct2.yaml
index fee5e7a06d6..63fa6d65e46 100644
--- a/examples/configs/sft_openmathinstruct2.yaml
+++ b/examples/configs/sft_openmathinstruct2.yaml
@@ -7,6 +7,7 @@ sft:
   val_global_batch_size: 128
   val_micro_batch_size: 2
   val_at_start: true
+  val_at_end: false
   seed: 42
 
 checkpointing:
diff --git a/examples/configs/sft_openmathinstruct2_megatron.yaml b/examples/configs/sft_openmathinstruct2_megatron.yaml
index 18aba3597bc..faca12e0aea 100644
--- a/examples/configs/sft_openmathinstruct2_megatron.yaml
+++ b/examples/configs/sft_openmathinstruct2_megatron.yaml
@@ -9,6 +9,7 @@ sft:
   val_global_batch_size: 128
   val_micro_batch_size: 1
   val_at_start: true
+  val_at_end: false
   seed: 42
 
 checkpointing:
diff --git a/examples/configs/vlm_grpo_3B.yaml b/examples/configs/vlm_grpo_3B.yaml
index 9eb3ea9d591..f9612007a4c 100644
--- a/examples/configs/vlm_grpo_3B.yaml
+++ b/examples/configs/vlm_grpo_3B.yaml
@@ -10,6 +10,7 @@ grpo:
   use_leave_one_out_baseline: true
   val_period: 10
   val_at_start: false
+  val_at_end: false
   overlong_filtering: false
   max_val_samples: 256
   val_batch_size: 256
diff --git a/examples/configs/vlm_grpo_3B_megatron.yaml b/examples/configs/vlm_grpo_3B_megatron.yaml
index 8bb4fb30e17..b32cd7df044 100644
--- a/examples/configs/vlm_grpo_3B_megatron.yaml
+++ b/examples/configs/vlm_grpo_3B_megatron.yaml
@@ -8,6 +8,7 @@ grpo:
   use_leave_one_out_baseline: true
   val_period: 10
   val_at_start: false
+  val_at_end: false
   overlong_filtering: false
   max_val_samples: 256
   val_batch_size: 256
diff --git a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml
index ae3740850dd..f3e3dcccc84 100644
--- a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml
+++ b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml
@@ -8,6 +8,7 @@ grpo:
   use_leave_one_out_baseline: true
   val_period: 10
   val_at_start: true
+  val_at_end: false
   overlong_filtering: false
   max_val_samples: null   # inferred from size of val dataset. for multi evals, repeat val ds via `num_repeats` in `ng_prepare_data`.
   val_batch_size: null
diff --git a/nemo_rl/algorithms/distillation.py b/nemo_rl/algorithms/distillation.py
index 5ec8671dd0f..6fa9689d1a0 100644
--- a/nemo_rl/algorithms/distillation.py
+++ b/nemo_rl/algorithms/distillation.py
@@ -80,6 +80,9 @@ class DistillationConfig(TypedDict):
     val_batch_size: int
     val_period: int
     val_at_start: bool
+    # Whether to run validation on the last training step. Setting this to True ensures the
+    # final checkpoint has validation metrics, which is required for get_best_checkpoint_path().
+    val_at_end: bool
     max_val_samples: int
     topk_logits_k: int
     seed: int
@@ -257,7 +260,11 @@ def setup(
     # Load validation dataset if provided
     val_dataloader: Optional[StatefulDataLoader] = None
     # If validation is enabled, load the validation dataloader
-    if distillation_config["val_period"] > 0 or distillation_config["val_at_start"]:
+    if (
+        distillation_config["val_period"] > 0
+        or distillation_config["val_at_start"]
+        or distillation_config["val_at_end"]
+    ):
         assert val_dataset is not None, (
             "Validation dataset is required if validation is enabled"
         )
@@ -539,6 +546,7 @@ def distillation_train(
     total_valid_tokens = distillation_save_state["total_valid_tokens"]
     val_period = master_config["distillation"]["val_period"]
     val_at_start = master_config["distillation"]["val_at_start"]
+    val_at_end = master_config["distillation"]["val_at_end"]
     colocated_inference = master_config["policy"]["generation"]["colocated"]["enabled"]
     max_epochs = master_config["distillation"][
         "max_num_epochs"
@@ -721,8 +729,10 @@ def distillation_train(
                     and (current_step + 1 == len(dataloader))
                 )
 
-                # Run validation if it's a validation step
-                if val_period > 0 and (total_steps + 1) % val_period == 0:
+                # Run validation if it's a validation step or last step with val_at_end
+                if (val_period > 0 and (total_steps + 1) % val_period == 0) or (
+                    val_at_end and is_last_step
+                ):
                     if NEED_REFIT and POLICY_GENERATION_STALE:
                         refit_policy_generation(
                             student_policy, student_generation, colocated_inference
diff --git a/nemo_rl/algorithms/dpo.py b/nemo_rl/algorithms/dpo.py
index 5ff865c87d8..dc532f546dc 100644
--- a/nemo_rl/algorithms/dpo.py
+++ b/nemo_rl/algorithms/dpo.py
@@ -66,6 +66,9 @@ class DPOConfig(TypedDict):
     val_global_batch_size: int
     val_micro_batch_size: int
     val_at_start: bool
+    # Whether to run validation on the last training step. Setting this to True ensures the
+    # final checkpoint has validation metrics, which is required for get_best_checkpoint_path().
+    val_at_end: bool
     seed: int
 
     reference_policy_kl_penalty: float
@@ -524,6 +527,7 @@ def dpo_train(
     # Validation configuration
     val_period = dpo_config["val_period"]
     val_at_start = dpo_config["val_at_start"]
+    val_at_end = dpo_config["val_at_end"]
     max_num_epochs = dpo_config["max_num_epochs"]
 
     # Run validation at the start if configured
@@ -582,8 +586,10 @@ def dpo_train(
                     and current_step + 1 == len(train_dataloader)
                 )
 
-                # Run validation if it's a validation step
-                if val_period > 0 and (total_steps + 1) % val_period == 0:
+                # Run validation if it's a validation step or last step with val_at_end
+                if (val_period > 0 and (total_steps + 1) % val_period == 0) or (
+                    val_at_end and is_last_step
+                ):
                     validation_result = validate(
                         policy,
                         val_dataloader,
diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py
index 71589210ce7..75e84f36a12 100644
--- a/nemo_rl/algorithms/grpo.py
+++ b/nemo_rl/algorithms/grpo.py
@@ -125,6 +125,9 @@ class GRPOConfig(TypedDict):
     val_period: int
     val_batch_size: int
     val_at_start: bool
+    # Whether to run validation on the last training step. Setting this to True ensures the
+    # final checkpoint has validation metrics, which is required for get_best_checkpoint_path().
+    val_at_end: bool
     max_val_samples: int
     skip_reference_policy_logprobs_calculation: NotRequired[bool]
     seed: int
@@ -279,7 +282,11 @@ def setup(
     # Load validation dataset if provided
     val_dataloader: Optional[StatefulDataLoader] = None
     # If validation is enabled, load the validation dataloader
-    if grpo_config["val_period"] > 0 or grpo_config["val_at_start"]:
+    if (
+        grpo_config["val_period"] > 0
+        or grpo_config["val_at_start"]
+        or grpo_config["val_at_end"]
+    ):
         assert val_dataset is not None, (
             "Validation dataset is required if validation is enabled"
         )
@@ -1158,6 +1165,7 @@ def grpo_train(
         "total_valid_tokens", 0
     )  # total valid tokens processed across all epochs; default to 0 for backward compatibility with older checkpoints
     val_at_start = master_config["grpo"]["val_at_start"]
+    val_at_end = master_config["grpo"]["val_at_end"]
     val_period = master_config["grpo"]["val_period"]
     colocated_inference = master_config["policy"]["generation"]["colocated"]["enabled"]
 
@@ -1565,8 +1573,10 @@ def grpo_train(
                     and (current_step + 1 == len(dataloader))
                 )
 
-                # Run validation if it's a validation step
-                if val_period > 0 and (total_steps + 1) % val_period == 0:
+                # Run validation if it's a validation step or last step with val_at_end
+                if (val_period > 0 and (total_steps + 1) % val_period == 0) or (
+                    val_at_end and is_last_step
+                ):
                     memory_tracker.snapshot_start_of_stage("Validation", dir())
                     if NEED_REFIT and POLICY_GENERATION_STALE:
                         refit_policy_generation(
@@ -2136,6 +2146,7 @@ def async_grpo_train(
     )  # Default to 0 for backward compatibility with older checkpoints
     val_period = master_config["grpo"]["val_period"]
     val_at_start = master_config["grpo"]["val_at_start"]
+    val_at_end = master_config["grpo"]["val_at_end"]
     colocated_inference = master_config["policy"]["generation"]["colocated"]["enabled"]
 
     assert not colocated_inference, (
@@ -2575,7 +2586,10 @@ def async_grpo_train(
                 val_metrics, validation_timings = None, None
                 is_last_step = step + 1 == master_config["grpo"]["max_num_steps"]
 
-                if val_period > 0 and (step + 1) % val_period == 0:
+                # Run validation if it's a validation step or last step with val_at_end
+                if (val_period > 0 and (step + 1) % val_period == 0) or (
+                    val_at_end and is_last_step
+                ):
                     # Pause trajectory collection during validation to reduce memory pressure
                     trajectory_collector.pause.remote()
 
diff --git a/nemo_rl/algorithms/rm.py b/nemo_rl/algorithms/rm.py
index 6a61cf6a3e0..584c217eb7e 100644
--- a/nemo_rl/algorithms/rm.py
+++ b/nemo_rl/algorithms/rm.py
@@ -67,6 +67,9 @@ class RMConfig(TypedDict):
     val_global_batch_size: int
     val_micro_batch_size: int
     val_at_start: bool
+    # Whether to run validation on the last training step. Setting this to True ensures the
+    # final checkpoint has validation metrics, which is required for get_best_checkpoint_path().
+    val_at_end: bool
     seed: int
 
 
@@ -459,6 +462,7 @@ def rm_train(
     # Validation configuration
     val_period = rm_config["val_period"]
     val_at_start = rm_config["val_at_start"]
+    val_at_end = rm_config["val_at_end"]
     max_num_epochs = rm_config["max_num_epochs"]
 
     # Run validation at the start if configured
@@ -515,8 +519,10 @@ def rm_train(
                     and current_step + 1 == len(train_dataloader)
                 )
 
-                # Run validation if it's a validation step
-                if val_period > 0 and (total_steps + 1) % val_period == 0:
+                # Run validation if it's a validation step or last step with val_at_end
+                if (val_period > 0 and (total_steps + 1) % val_period == 0) or (
+                    val_at_end and is_last_step
+                ):
                     val_metrics, validation_timings = validate(
                         policy,
                         val_dataloader,
diff --git a/nemo_rl/algorithms/sft.py b/nemo_rl/algorithms/sft.py
index 130cd2fdd62..dcd7b9d0257 100644
--- a/nemo_rl/algorithms/sft.py
+++ b/nemo_rl/algorithms/sft.py
@@ -70,6 +70,9 @@ class SFTConfig(TypedDict):
     val_global_batch_size: int
     val_micro_batch_size: int
     val_at_start: bool
+    # Whether to run validation on the last training step. Setting this to True ensures the
+    # final checkpoint has validation metrics, which is required for get_best_checkpoint_path().
+    val_at_end: bool
     seed: int
 
 
@@ -385,6 +388,7 @@ def sft_train(
     # Validation configuration
     val_period = sft_config["val_period"]
     val_at_start = sft_config["val_at_start"]
+    val_at_end = sft_config["val_at_end"]
     max_num_epochs = sft_config["max_num_epochs"]
 
     # Run validation at the start if configured
@@ -465,8 +469,10 @@ def sft_train(
                     and current_step + 1 == len(train_dataloader)
                 )
 
-                # Run validation if it's a validation step
-                if val_period > 0 and (total_steps + 1) % val_period == 0:
+                # Run validation if it's a validation step or last step with val_at_end
+                if (val_period > 0 and (total_steps + 1) % val_period == 0) or (
+                    val_at_end and is_last_step
+                ):
                     val_metrics, validation_timings = validate(
                         policy,
                         val_dataloader,
diff --git a/nemo_rl/utils/checkpoint.py b/nemo_rl/utils/checkpoint.py
index 05e0ee2f3a7..740e4cfc2d0 100644
--- a/nemo_rl/utils/checkpoint.py
+++ b/nemo_rl/utils/checkpoint.py
@@ -230,25 +230,41 @@ def get_best_checkpoint_path(self) -> Optional[str]:
         """Get the path to the best checkpoint based on the metric.
 
         Returns the path to the checkpoint with the best metric value. If no checkpoints
-        exist, returns None. If the metric isn't found, we warn and return the latest checkpoint.
+        exist, returns None. If some checkpoints are missing the metric, they are filtered
+        out with a warning. If no checkpoints have the metric, returns the latest checkpoint.
 
         Returns:
-            Optional[str]: Path to the best checkpoint, or None if no valid checkpoints exist.
+            Optional[str]: Path to the best checkpoint, or None if no checkpoints exist.
         """
         checkpoint_history = _load_checkpoint_history(self.checkpoint_dir)
         if len(checkpoint_history) == 0:
             return None
-        # sort by metric value
-        if self.metric_name not in checkpoint_history[0][2]:
+
+        # Filter checkpoints that have the metric
+        valid_checkpoints = [c for c in checkpoint_history if self.metric_name in c[2]]
+        ignored_count = len(checkpoint_history) - len(valid_checkpoints)
+
+        if ignored_count > 0:
+            ignored_steps = [
+                c[0] for c in checkpoint_history if self.metric_name not in c[2]
+            ]
+            warnings.warn(
+                f"Ignoring {ignored_count} checkpoint(s) at step(s) {ignored_steps} that do not have "
+                f"metric '{self.metric_name}'. Consider enabling val_at_end or adjusting val_period "
+                f"to align with max_steps."
+            )
+
+        if len(valid_checkpoints) == 0:
             warnings.warn(
-                f"Metric {self.metric_name} not found in checkpoint history. Returning last"
+                f"No checkpoints contain metric '{self.metric_name}'. Returning latest checkpoint. "
+                f"Consider enabling val_at_end or adjusting val_period to align with max_steps."
             )
             return self.get_latest_checkpoint_path()
 
-        checkpoint_history.sort(
+        valid_checkpoints.sort(
             key=lambda x: x[2][self.metric_name], reverse=self.higher_is_better
         )
-        return str(checkpoint_history[0][1])
+        return str(valid_checkpoints[0][1])
 
     def get_latest_checkpoint_path(self) -> Optional[str]:
         """Get the path to the latest checkpoint.
diff --git a/research/template_project/configs/grpo_math_1B.yaml b/research/template_project/configs/grpo_math_1B.yaml
index cd6875c1d6e..c28e86d1573 100644
--- a/research/template_project/configs/grpo_math_1B.yaml
+++ b/research/template_project/configs/grpo_math_1B.yaml
@@ -9,6 +9,7 @@ grpo:
   use_leave_one_out_baseline: true
   val_period: 10
   val_at_start: false
+  val_at_end: false
   overlong_filtering: false
   max_val_samples: 256
   val_batch_size: 256
diff --git a/tests/functional/test_converter_roundtrip.py b/tests/functional/test_converter_roundtrip.py
index 7e6bf8d4a78..afe6ef06592 100644
--- a/tests/functional/test_converter_roundtrip.py
+++ b/tests/functional/test_converter_roundtrip.py
@@ -52,6 +52,7 @@ def create_test_config() -> Dict[str, Any]:
             "val_global_batch_size": 4,
             "val_micro_batch_size": 2,
             "val_at_start": False,
+            "val_at_end": False,
             "seed": 42,
         },
         "checkpointing": {
diff --git a/tests/unit/algorithms/test_distillation.py b/tests/unit/algorithms/test_distillation.py
index ec8ed37f635..a0dfc19d695 100644
--- a/tests/unit/algorithms/test_distillation.py
+++ b/tests/unit/algorithms/test_distillation.py
@@ -128,6 +128,7 @@ def val_iter(self):
             "val_period": 100,
             "val_batch_size": 1,
             "val_at_start": False,
+            "val_at_end": False,
             "max_val_samples": 10,
             "topk_logits_k": 64,
             "num_prompts_per_step": 1,
@@ -444,6 +445,7 @@ def test_noncolocated_inference_requires_explicit_gpus_per_node_single_node():
             "num_prompts_per_step": 1,  # Config extraction requires this key
             "val_period": 0,  # Config extraction requires this key
             "val_at_start": False,  # Config extraction requires this key
+            "val_at_end": False,  # Config extraction requires this key
         },
         "data": {"shuffle": False},
         "logger": {},  # Config extraction requires this key
@@ -516,6 +518,7 @@ def test_distillation_setup_non_colocated_smoke(monkeypatch):
             "max_num_steps": 100,
             "val_period": 0,
             "val_at_start": False,
+            "val_at_end": False,
         },
         "data": {"shuffle": False},
         "logger": {},
@@ -624,6 +627,7 @@ def test_noncolocated_inference_requires_explicit_gpus_per_node_multi_node():
             "num_prompts_per_step": 1,  # Config extraction requires this key
             "val_period": 0,  # Config extraction requires this key
             "val_at_start": False,  # Config extraction requires this key
+            "val_at_end": False,  # Config extraction requires this key
         },
         "data": {"shuffle": False},
         "logger": {},  # Config extraction requires this key
diff --git a/tests/unit/algorithms/test_dpo.py b/tests/unit/algorithms/test_dpo.py
index 45d59fe9902..b2155ac91f3 100644
--- a/tests/unit/algorithms/test_dpo.py
+++ b/tests/unit/algorithms/test_dpo.py
@@ -183,6 +183,7 @@ def val_iter(self):
             "val_global_batch_size": 1,
             "val_micro_batch_size": 1,
             "val_at_start": False,
+            "val_at_end": False,
         },
         "policy": {
             "train_global_batch_size": 2,
diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py
index c19d3153dc1..1729f4eac1e 100644
--- a/tests/unit/algorithms/test_grpo.py
+++ b/tests/unit/algorithms/test_grpo.py
@@ -724,6 +724,7 @@ def test_noncolocated_inference_requires_explicit_gpus_per_node_single_node():
             "num_prompts_per_step": 1,
             "val_period": 0,
             "val_at_start": False,
+            "val_at_end": False,
             "use_dynamic_sampling": False,
             "batch_multiplier": 1,
         },
@@ -782,6 +783,7 @@ def test_noncolocated_inference_requires_explicit_gpus_per_node_multi_node():
             "num_prompts_per_step": 1,
             "val_period": 0,
             "val_at_start": False,
+            "val_at_end": False,
             "use_dynamic_sampling": False,
             "batch_multiplier": 1,
         },
@@ -939,6 +941,7 @@ def init_collective(self, *_args, **_kwargs):
             "val_period": 0,
             "val_batch_size": 1,
             "val_at_start": False,
+            "val_at_end": False,
             "max_val_samples": 1,
             "use_dynamic_sampling": False,
             "batch_multiplier": 1,
@@ -1246,6 +1249,7 @@ def val_iter(self):
             "val_period": 100,
             "val_batch_size": 1,
             "val_at_start": False,
+            "val_at_end": False,
             "max_val_samples": 10,
             "seed": 42,
             "advantage_normalization": "global",
diff --git a/tests/unit/algorithms/test_rm.py b/tests/unit/algorithms/test_rm.py
index 8dabbedcb39..0c369d028c5 100644
--- a/tests/unit/algorithms/test_rm.py
+++ b/tests/unit/algorithms/test_rm.py
@@ -90,6 +90,7 @@ def val_iter(self):
             "val_global_batch_size": 1,
             "val_micro_batch_size": 1,
             "val_at_start": False,
+            "val_at_end": False,
         },
         "policy": {
             "train_global_batch_size": 1,
diff --git a/tests/unit/algorithms/test_sft.py b/tests/unit/algorithms/test_sft.py
index fffd06187db..c507f8a9870 100644
--- a/tests/unit/algorithms/test_sft.py
+++ b/tests/unit/algorithms/test_sft.py
@@ -72,6 +72,7 @@ def val_iter(self):
             "val_global_batch_size": 1,
             "val_micro_batch_size": 1,
             "val_at_start": False,
+            "val_at_end": False,
         },
         "policy": {
             "train_global_batch_size": 1,
diff --git a/tests/unit/utils/test_checkpoint.py b/tests/unit/utils/test_checkpoint.py
index 002524cc716..9d42fa7d18e 100644
--- a/tests/unit/utils/test_checkpoint.py
+++ b/tests/unit/utils/test_checkpoint.py
@@ -347,3 +347,123 @@ def test_get_latest_checkpoint_path_across_digits(checkpoint_manager, checkpoint
     )
     assert len(all_checkpoints) == checkpoint_manager.keep_top_k
     assert all_checkpoints[-1].name == f"step_{max(steps)}"
+
+
+def test_get_best_checkpoint_path_no_checkpoints(checkpoint_manager, checkpoint_dir):
+    """Test that get_best_checkpoint_path returns None when no checkpoints exist."""
+    result = checkpoint_manager.get_best_checkpoint_path()
+    assert result is None
+
+
+def test_get_best_checkpoint_path_some_missing_metric(tmp_path):
+    """Test that get_best_checkpoint_path filters out checkpoints missing the metric and warns."""
+    # Use keep_top_k=None to keep all checkpoints for this test
+    config = {
+        "enabled": True,
+        "checkpoint_dir": str((tmp_path.resolve() / "checkpoints")),
+        "metric_name": "loss",
+        "higher_is_better": False,
+        "keep_top_k": None,  # Keep all checkpoints
+    }
+    manager = CheckpointManager(config)
+
+    # Create checkpoints where some have the metric and others don't
+    steps = [1, 2, 3, 4]
+    training_infos = [
+        {"loss": 0.5},  # step 1 - has loss
+        {"other_metric": 0.8},  # step 2 - missing loss
+        {"loss": 0.3},  # step 3 - has loss (best)
+        {},  # step 4 - missing loss
+    ]
+
+    for step, training_info in zip(steps, training_infos):
+        tmp_dir = manager.init_tmp_checkpoint(step, training_info)
+        manager.finalize_checkpoint(tmp_dir)
+
+    # Should warn about missing metrics but still return the best checkpoint
+    import warnings
+
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        best_path = manager.get_best_checkpoint_path()
+
+        # Should have warned about 2 checkpoints missing the metric
+        assert len(w) == 1
+        assert "Ignoring 2 checkpoint(s)" in str(w[0].message)
+        assert "val_at_end" in str(w[0].message)
+
+    # Should return the checkpoint with the best (lowest) loss
+    with open(Path(best_path) / "training_info.json", "r") as f:
+        metadata = json.load(f)
+        assert metadata["loss"] == 0.3  # step 3 has the best loss
+
+
+def test_get_best_checkpoint_path_all_missing_metric(tmp_path):
+    """Test that get_best_checkpoint_path returns latest checkpoint when all are missing the metric."""
+    # Use keep_top_k=None to keep all checkpoints for this test
+    config = {
+        "enabled": True,
+        "checkpoint_dir": str((tmp_path.resolve() / "checkpoints")),
+        "metric_name": "loss",
+        "higher_is_better": False,
+        "keep_top_k": None,  # Keep all checkpoints
+    }
+    manager = CheckpointManager(config)
+
+    # Create checkpoints where none have the required metric
+    steps = [1, 2, 3]
+    training_infos = [
+        {"other_metric": 0.5},  # step 1 - missing loss
+        {},  # step 2 - missing loss
+        {"different_metric": 0.3},  # step 3 - missing loss
+    ]
+
+    for step, training_info in zip(steps, training_infos):
+        tmp_dir = manager.init_tmp_checkpoint(step, training_info)
+        manager.finalize_checkpoint(tmp_dir)
+
+    # Should warn and return latest checkpoint when no checkpoints have the metric
+    import warnings
+
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        best_path = manager.get_best_checkpoint_path()
+
+        # Should have warned twice: once about ignoring all checkpoints, once about returning latest
+        assert len(w) == 2
+        assert "Ignoring 3 checkpoint(s)" in str(w[0].message)
+        assert "No checkpoints contain metric 'loss'" in str(w[1].message)
+        assert "Returning latest checkpoint" in str(w[1].message)
+        assert "val_at_end" in str(w[1].message)
+
+    # Should return the latest checkpoint (step 3)
+    assert Path(best_path).name == "step_3"
+
+
+def test_get_best_checkpoint_path_higher_is_better(tmp_path):
+    """Test get_best_checkpoint_path with higher_is_better=True."""
+    config = {
+        "enabled": True,
+        "checkpoint_dir": str((tmp_path.resolve() / "checkpoints")),
+        "metric_name": "accuracy",
+        "higher_is_better": True,
+        "keep_top_k": None,  # Keep all
+    }
+    manager = CheckpointManager(config)
+
+    # Create checkpoints with different accuracy values
+    steps = [1, 2, 3]
+    accuracies = [0.7, 0.9, 0.8]  # step 2 has the best accuracy
+
+    for step, acc in zip(steps, accuracies):
+        training_info = {"accuracy": acc}
+        tmp_dir = manager.init_tmp_checkpoint(step, training_info)
+        manager.finalize_checkpoint(tmp_dir)
+
+    # Get best checkpoint path
+    best_path = manager.get_best_checkpoint_path()
+
+    # Verify it's the checkpoint with highest accuracy
+    with open(Path(best_path) / "training_info.json", "r") as f:
+        metadata = json.load(f)
+        assert metadata["accuracy"] == 0.9  # step 2