add for rest of policies

src/lerobot/policies/pi0/modeling_pi0.py

@@ -995,7 +995,14 @@ class PI0Policy(PreTrainedPolicy):
 
         # Initialize model without loading weights
         # Check if dataset_stats were provided in kwargs
-        model = cls(config, **kwargs)
+        if _transformers_available:
+            from transformers.modeling_utils import no_init_weights
+
+            with no_init_weights():
+                model = cls(config, **kwargs)
+            model.model.paligemma_with_expert.paligemma.tie_weights()
+        else:
+            model = cls(config, **kwargs)
 
         # Now manually load and remap the state dict
         try:

src/lerobot/policies/pi05/modeling_pi05.py

@@ -969,6 +969,7 @@ class PI05Policy(PreTrainedPolicy):
         # Check if dataset_stats were provided in kwargs
         if _transformers_available:
             from transformers.modeling_utils import no_init_weights
+
             with no_init_weights():
                 model = cls(config, **kwargs)
             model.model.paligemma_with_expert.paligemma.tie_weights()

src/lerobot/policies/pi0_fast/modeling_pi0_fast.py

@@ -895,7 +895,14 @@ class PI0FastPolicy(PreTrainedPolicy):
 
         # Initialize model without loading weights
         # Check if dataset_stats were provided in kwargs
-        model = cls(config, **kwargs)
+        if _transformers_available:
+            from transformers.modeling_utils import no_init_weights
+
+            with no_init_weights():
+                model = cls(config, **kwargs)
+            model.model.paligemma_with_expert.paligemma.tie_weights()
+        else:
+            model = cls(config, **kwargs)
 
         # Now manually load and remap the state dict
         try: