Add optional batch size in option args for replay policy (#504)

* add extra batch size in option for replay

* mark obs optional

Author: xyao-nv

gr00t/policy/replay_policy.py

@@ -300,16 +300,24 @@ def _get_action(
         """Replay the next action chunk from the dataset.
 
         Args:
-            observation: Batched observation dictionary (used for validation, not inference)
-            options: Optional parameters (currently unused)
+            observation: Optional batched observation dictionary (used for validation, not inference)
+            options: Optional parameters
+                - batch_size: int - Batch size to use for the action chunk
 
         Returns:
             Tuple of (actions_dict, info_dict) where actions_dict contains action chunks
             with shape (B, action_horizon, D) for each action key
         """
         # Infer batch size from observation
-        first_video_key = self.modality_configs["video"].modality_keys[0]
-        batch_size = observation["video"][first_video_key].shape[0]
+        if observation is not None:
+            first_video_key = self.modality_configs["video"].modality_keys[0]
+            batch_size = observation["video"][first_video_key].shape[0]
+        # If batch size is not provided in observation, check if it's provided in options
+        elif "batch_size" in options:
+            batch_size = options["batch_size"]
+        else:
+            batch_size = 1
+            print("No batch size provided, using default batch size of 1")
         # Note that this can differ form the execution horizon, as the policy can predict more steps than what's actually executed.
         action_horizon = (
             self.modality_configs["action"].delta_indices[-1]