diff --git a/src/lerobot/policies/g05/configuration_g05.py b/src/lerobot/policies/g05/configuration_g05.py index e13aa04ea..ec8b7ce83 100644 --- a/src/lerobot/policies/g05/configuration_g05.py +++ b/src/lerobot/policies/g05/configuration_g05.py @@ -168,6 +168,7 @@ class G05Config(PreTrainedConfig): discrete_action: bool = True continuous_action: bool = False return_continuous_action: bool = False + model_weights_to_bf16: bool = True policy_action_dim: int = 20 policy_state_dim: int = 20 diff --git a/src/lerobot/policies/g05/modeling_g05.py b/src/lerobot/policies/g05/modeling_g05.py index 586813805..b9381bc6e 100644 --- a/src/lerobot/policies/g05/modeling_g05.py +++ b/src/lerobot/policies/g05/modeling_g05.py @@ -165,10 +165,25 @@ class G05Policy(PreTrainedPolicy): if callable(reset): reset() + def _apply_author_inference_precision(self) -> None: + """Match the released serving path's BF16 weights with declared FP32 islands.""" + + self.backend.to(dtype=torch.bfloat16) + apply_fp32_params = getattr(self.backend, "apply_fp32_params", None) + if callable(apply_fp32_params): + apply_fp32_params() + def to(self, *args, **kwargs) -> G05Policy: - """Move the author ActionCodec sidecar along with the policy module.""" + """Apply author inference precision and move the ActionCodec sidecar.""" result = super().to(*args, **kwargs) + explicit_dtype = "dtype" in kwargs or any(isinstance(arg, torch.dtype | Tensor) for arg in args) + if ( + self.config.model_weights_to_bf16 + and not explicit_dtype + and next(self.backend.parameters()).device.type == "cuda" + ): + self._apply_author_inference_precision() action_tokenizer = getattr(self.backend, "action_tokenizer", None) move_tokenizer = getattr(action_tokenizer, "to", None) if callable(move_tokenizer): diff --git a/tests/policies/g05/test_g05.py b/tests/policies/g05/test_g05.py index a3a599e14..866e92f11 100644 --- a/tests/policies/g05/test_g05.py +++ b/tests/policies/g05/test_g05.py @@ -685,6 +685,25 @@ def test_policy_to_moves_non_module_action_tokenizer_sidecar(): assert backend.action_tokenizer.device == next(policy.parameters()).device +def test_author_inference_precision_preserves_declared_fp32_parameters(): + class MixedPrecisionBackend(TinyG05Backend): + def __init__(self): + super().__init__() + self.bulk_weight = nn.Parameter(torch.ones(2)) + self.precision_weight = nn.Parameter(torch.ones(2)) + + def apply_fp32_params(self): + self.precision_weight.data = self.precision_weight.data.float() + + backend = MixedPrecisionBackend() + policy = G05Policy(_config(), backend=backend) + + policy._apply_author_inference_precision() + + assert backend.bulk_weight.dtype is torch.bfloat16 + assert backend.precision_weight.dtype is torch.float32 + + def test_batch_two_preserves_each_raw_task_and_every_camera_slot(): backend = TinyG05Backend() policy = G05Policy(_config(), backend=backend)