add different dtype support

src/lerobot/policies/xvla/configuration_xvla.py

@@ -51,6 +51,7 @@ class XVLAConfig(PreTrainedConfig):
     n_obs_steps: int = 1
     chunk_size: int = 32
     n_action_steps: int = 32
+    dtype: str = "float32"  # Options: "bfloat16", "float32"
 
     normalization_mapping: dict[str, NormalizationMode] = field(
         default_factory=lambda: {
@@ -119,6 +120,8 @@ class XVLAConfig(PreTrainedConfig):
             )
         if self.num_image_views is not None and self.num_image_views <= 0:
             raise ValueError("`num_image_views` must be > 0 when specified.")
+        if self.dtype not in ["bfloat16", "float32"]:
+            raise ValueError(f"Invalid dtype: {self.dtype}")
         self._florence_config_obj: Florence2Config | None = None
 
     def get_florence_config(self) -> Florence2Config:

src/lerobot/policies/xvla/modeling_xvla.py

@@ -88,6 +88,22 @@ class XVLAModel(nn.Module):
         # Apply freezing based on config
         self._apply_freezing()
 
+        # Apply dtype casting based on config
+        self._apply_dtype()
+
+    def _get_target_dtype(self) -> torch.dtype:
+        """Get the target dtype based on config."""
+        if self.config.dtype == "bfloat16":
+            return torch.bfloat16
+        return torch.float32
+
+    def _apply_dtype(self) -> None:
+        """
+        Apply dtype casting to model components based on config.
+        """
+        target_dtype = self._get_target_dtype()
+        self.to(dtype=target_dtype)
+
     def _apply_freezing(self) -> None:
         """
         Freeze VLM vision and language encoders based on config options.
@@ -451,6 +467,8 @@ class XVLAPolicy(PreTrainedPolicy):
         instance.load_state_dict(state_dict, strict=True)
         print("Loaded XVLA checkpoint")
         # step 5: finalize
+        # Reapply dtype after loading state dict
+        instance.model._apply_dtype()
         instance.to(config.device)
         instance.eval()
         return instance