Merge pull request #274 from mananroongta/manan/agent_ppo_trainer

Fix : Prevent KeyError in _pad_dataproto_to_world_size

Author: jeffreysijuntan

rllm/trainer/verl/agent_ppo_trainer.py

@@ -1040,8 +1040,10 @@ def _pad_dataproto_to_world_size(self, batch):
         # for the padded dataproto, make the traj mask to 0. is_last_step also False
         for i in range(pad_size):
             idx = original_batch_size + i
-            batch.non_tensor_batch["is_last_step"][idx] = False
-            batch.non_tensor_batch["is_pad_step"][idx] = True
+            if "is_last_step" in batch.non_tensor_batch:
+                batch.non_tensor_batch["is_last_step"][idx] = False
+            if "is_pad_step" in batch.non_tensor_batch:
+                batch.non_tensor_batch["is_pad_step"][idx] = True
 
         return batch