diff --git a/src/lerobot/policies/xvla/configuration_xvla.py b/src/lerobot/policies/xvla/configuration_xvla.py index c86e921bf..e4a2c2a74 100644 --- a/src/lerobot/policies/xvla/configuration_xvla.py +++ b/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: diff --git a/src/lerobot/policies/xvla/modeling_xvla.py b/src/lerobot/policies/xvla/modeling_xvla.py index 8da58264f..bc691ac04 100644 --- a/src/lerobot/policies/xvla/modeling_xvla.py +++ b/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