fix CI bugs

Author: hijkzzz

nemo_rl/algorithms/loss_functions.py

@@ -46,9 +46,9 @@ class ClippedPGLossConfig(TypedDict):
     use_importance_sampling_correction: bool
     truncated_importance_sampling_ratio: float | None
     # Type of truncated importance sampling: "tis" (clamp max) or "icepop" (filter [min, max])
-    truncated_importance_sampling_type: NotRequired[str]
+    truncated_importance_sampling_type: NotRequired[str | None]
     # Lower bound for ICE-POP filtering (default 0.5)
-    truncated_importance_sampling_ratio_min: NotRequired[float]
+    truncated_importance_sampling_ratio_min: NotRequired[float | None]
     token_level_loss: bool
     # If True, apply the off-policy importance-sampling correction at the
     # sequence level (one weight per generated sample), as in GSPO.

nemo_rl/algorithms/reward_functions.py

@@ -47,7 +47,7 @@ class RewardShapingConfig(TypedDict):
     # Stop properly penalty: scale factor for rewards of truncated responses (0-1).
     # When set to 0, truncated responses get zero reward.
     # When set to 1, no penalty is applied (default behavior).
-    stop_properly_penalty_coef: NotRequired[float]
+    stop_properly_penalty_coef: NotRequired[float | None]
 
 
 def apply_reward_shaping(

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -543,6 +543,7 @@ def generate(
                     "logprobs": torch.zeros((0, 0), dtype=torch.float),
                     "generation_lengths": torch.zeros(0, dtype=torch.long),
                     "unpadded_sequence_lengths": torch.zeros(0, dtype=torch.long),
+                    "truncated": torch.zeros(0, dtype=torch.bool),
                 }
             )
 

nemo_rl/models/generation/vllm/vllm_worker_async.py

@@ -727,12 +727,18 @@ async def process_single_sample(sample_idx):
                     device=input_ids_single_row.device,
                 )
 
+                # Not truncated since no generation was attempted (length constraint)
+                truncated_tensor = torch.tensor(
+                    [False], dtype=torch.bool, device=input_ids_single_row.device
+                )
+
                 result_batch = BatchedDataDict[GenerationOutputSpec](
                     {
                         "output_ids": output_ids_single_item_batched,
                         "logprobs": logprobs_single_item,
                         "generation_lengths": generation_lengths_tensor,
                         "unpadded_sequence_lengths": unpadded_sequence_lengths_tensor,
+                        "truncated": truncated_tensor,
                     }
                 )