From 28fa7eae729702c142d2e805cdd97af85243e543 Mon Sep 17 00:00:00 2001 From: AdilZouitine Date: Tue, 23 Sep 2025 10:03:39 +0200 Subject: [PATCH] refactor(pi0): rename PI0OpenPIConfig and PI0OpenPIPolicy to PI0Config and PI0Policy - Updated imports and references throughout the codebase to reflect the new naming convention. - Introduced a new processor file for PI0 to handle pre-processing and post-processing steps. - Adjusted tests to utilize the renamed classes, ensuring consistency and functionality. - Enhanced clarity and maintainability by removing outdated naming conventions. --- src/lerobot/policies/__init__.py | 4 ++-- src/lerobot/policies/factory.py | 14 +++++++------- src/lerobot/policies/pi0/__init__.py | 8 ++++---- .../policies/pi0/configuration_pi0openpi.py | 2 +- src/lerobot/policies/pi0/modeling_pi0openpi.py | 10 +++++----- .../{pi05 => pi0}/processor_pi0_openpi.py | 12 ++++++------ .../policies/pi05/modeling_pi05openpi.py | 2 +- .../{pi0 => pi05}/processor_pi05openpi.py | 4 ++-- tests/policies/pi0_pi05/test_pi0_openpi.py | 16 +++++++--------- .../pi0_pi05/test_pi0_original_vs_lerobot.py | 18 +++++++----------- 10 files changed, 42 insertions(+), 48 deletions(-) rename src/lerobot/policies/{pi05 => pi0}/processor_pi0_openpi.py (94%) rename src/lerobot/policies/{pi0 => pi05}/processor_pi05openpi.py (97%) diff --git a/src/lerobot/policies/__init__.py b/src/lerobot/policies/__init__.py index c0b12c121..d0005d124 100644 --- a/src/lerobot/policies/__init__.py +++ b/src/lerobot/policies/__init__.py @@ -14,7 +14,7 @@ from .act.configuration_act import ACTConfig as ACTConfig from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig -from .pi0.configuration_pi0openpi import PI0OpenPIConfig as PI0OpenPIConfig +from .pi0.configuration_pi0openpi import PI0Config from .pi05.configuration_pi05openpi import PI05OpenPIConfig as PI05OpenPIConfig from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig from .smolvla.processor_smolvla import SmolVLANewLineProcessor @@ -24,7 +24,7 @@ from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig __all__ = [ "ACTConfig", "DiffusionConfig", - "PI0OpenPIConfig", + "PI0Config", "PI05OpenPIConfig", "SmolVLAConfig", "TDMPCConfig", diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 2e04aad82..5f669ef73 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -31,7 +31,7 @@ from lerobot.envs.configs import EnvConfig from lerobot.envs.utils import env_to_policy_features from lerobot.policies.act.configuration_act import ACTConfig from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig -from lerobot.policies.pi0.configuration_pi0openpi import PI0OpenPIConfig +from lerobot.policies.pi0.configuration_pi0openpi import PI0Config from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig from lerobot.policies.pi05.configuration_pi05openpi import PI05OpenPIConfig from lerobot.policies.pretrained import PreTrainedPolicy @@ -140,7 +140,7 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig: elif policy_type == "pi0fast": return PI0FASTConfig(**kwargs) elif policy_type == "pi0": - return PI0OpenPIConfig(**kwargs) + return PI0Config(**kwargs) elif policy_type == "pi05": return PI05OpenPIConfig(**kwargs) elif policy_type == "sac": @@ -150,7 +150,7 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig: elif policy_type == "reward_classifier": return RewardClassifierConfig(**kwargs) elif policy_type == "pi0_openpi": - return PI0OpenPIConfig(**kwargs) + return PI0Config(**kwargs) elif policy_type == "pi05_openpi": return PI05OpenPIConfig(**kwargs) else: @@ -272,16 +272,16 @@ def make_pre_post_processors( dataset_stats=kwargs.get("dataset_stats"), ) - elif isinstance(policy_cfg, PI0OpenPIConfig): - from lerobot.policies.pi0_openpi.processor_pi0_openpi import make_pi0_openpi_pre_post_processors + elif isinstance(policy_cfg, PI0Config): + from lerobot.policies.pi0.processor_pi0_openpi import make_pi0_pre_post_processors - processors = make_pi0_openpi_pre_post_processors( + processors = make_pi0_pre_post_processors( config=policy_cfg, dataset_stats=kwargs.get("dataset_stats"), ) elif isinstance(policy_cfg, PI05OpenPIConfig): - from lerobot.policies.pi05_openpi.processor_pi05openpi import make_pi05_openpi_pre_post_processors + from lerobot.policies.pi05.processor_pi05openpi import make_pi05_openpi_pre_post_processors processors = make_pi05_openpi_pre_post_processors( config=policy_cfg, diff --git a/src/lerobot/policies/pi0/__init__.py b/src/lerobot/policies/pi0/__init__.py index 449461f93..e17e458f3 100644 --- a/src/lerobot/policies/pi0/__init__.py +++ b/src/lerobot/policies/pi0/__init__.py @@ -14,8 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .configuration_pi0openpi import PI0OpenPIConfig -from .modeling_pi0openpi import PI0OpenPIPolicy -from .processor_pi0_openpi import make_pi0_openpi_pre_post_processors +from .configuration_pi0openpi import PI0Config +from .modeling_pi0openpi import PI0Policy +from .processor_pi0_openpi import make_pi0_pre_post_processors -__all__ = ["PI0OpenPIConfig", "PI0OpenPIPolicy", "make_pi0_openpi_pre_post_processors"] +__all__ = ["PI0Config", "PI0Policy", "make_pi0_pre_post_processors"] diff --git a/src/lerobot/policies/pi0/configuration_pi0openpi.py b/src/lerobot/policies/pi0/configuration_pi0openpi.py index 7402b2b89..b7f19eef6 100644 --- a/src/lerobot/policies/pi0/configuration_pi0openpi.py +++ b/src/lerobot/policies/pi0/configuration_pi0openpi.py @@ -24,7 +24,7 @@ from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig @PreTrainedConfig.register_subclass("pi0") @dataclass -class PI0OpenPIConfig(PreTrainedConfig): +class PI0Config(PreTrainedConfig): # Model architecture paligemma_variant: str = "gemma_2b" action_expert_variant: str = "gemma_300m" diff --git a/src/lerobot/policies/pi0/modeling_pi0openpi.py b/src/lerobot/policies/pi0/modeling_pi0openpi.py index 3b8c2851a..b57a13eab 100644 --- a/src/lerobot/policies/pi0/modeling_pi0openpi.py +++ b/src/lerobot/policies/pi0/modeling_pi0openpi.py @@ -31,7 +31,7 @@ from transformers.models.paligemma.modeling_paligemma import PaliGemmaForConditi from lerobot.configs.policies import PreTrainedConfig from lerobot.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE -from lerobot.policies.pi0_openpi.configuration_pi0openpi import PI0OpenPIConfig +from lerobot.policies.pi0.configuration_pi0openpi import PI0Config from lerobot.policies.pretrained import PreTrainedPolicy, T @@ -490,7 +490,7 @@ class PaliGemmaWithExpertModel( class PI0Pytorch(nn.Module): # see openpi `PI0Pytorch` """Core PI0 PyTorch model.""" - def __init__(self, config: PI0OpenPIConfig): + def __init__(self, config: PI0Config): super().__init__() self.config = config @@ -839,15 +839,15 @@ $(python -c "import transformers, os; print(os.path.dirname(transformers.__file_ return self.action_out_proj(suffix_out) -class PI0OpenPIPolicy(PreTrainedPolicy): +class PI0Policy(PreTrainedPolicy): """PI0 OpenPI Policy for LeRobot.""" - config_class = PI0OpenPIConfig + config_class = PI0Config name = "pi0" def __init__( # see lerobot pi0 `__init__` self, - config: PI0OpenPIConfig, + config: PI0Config, ): """ Args: diff --git a/src/lerobot/policies/pi05/processor_pi0_openpi.py b/src/lerobot/policies/pi0/processor_pi0_openpi.py similarity index 94% rename from src/lerobot/policies/pi05/processor_pi0_openpi.py rename to src/lerobot/policies/pi0/processor_pi0_openpi.py index 14f148d92..f311e023f 100644 --- a/src/lerobot/policies/pi05/processor_pi0_openpi.py +++ b/src/lerobot/policies/pi0/processor_pi0_openpi.py @@ -18,7 +18,7 @@ import torch from lerobot.configs.types import PipelineFeatureType, PolicyFeature from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME -from lerobot.policies.pi0_openpi.configuration_pi0openpi import PI0OpenPIConfig +from lerobot.policies.pi0.configuration_pi0openpi import PI0Config from lerobot.processor import ( AddBatchDimensionProcessorStep, ComplementaryDataProcessorStep, @@ -35,8 +35,8 @@ from lerobot.processor import ( from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action -@ProcessorStepRegistry.register(name="pi0_openpi_new_line_processor") -class Pi0OpenPINewLineProcessor(ComplementaryDataProcessorStep): +@ProcessorStepRegistry.register(name="pi0_new_line_processor") +class Pi0NewLineProcessor(ComplementaryDataProcessorStep): """ Ensures that the task description string ends with a newline character. @@ -92,8 +92,8 @@ class Pi0OpenPINewLineProcessor(ComplementaryDataProcessorStep): return features -def make_pi0_openpi_pre_post_processors( - config: PI0OpenPIConfig, +def make_pi0_pre_post_processors( + config: PI0Config, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, ) -> tuple[ PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], @@ -128,7 +128,7 @@ def make_pi0_openpi_pre_post_processors( input_steps: list[ProcessorStep] = [ RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one AddBatchDimensionProcessorStep(), - Pi0OpenPINewLineProcessor(), # Add newlines before tokenization for PaliGemma + Pi0NewLineProcessor(), # Add newlines before tokenization for PaliGemma TokenizerProcessorStep( tokenizer_name="google/paligemma-3b-pt-224", max_length=config.tokenizer_max_length, diff --git a/src/lerobot/policies/pi05/modeling_pi05openpi.py b/src/lerobot/policies/pi05/modeling_pi05openpi.py index e4043e6ca..22af08b35 100644 --- a/src/lerobot/policies/pi05/modeling_pi05openpi.py +++ b/src/lerobot/policies/pi05/modeling_pi05openpi.py @@ -31,7 +31,7 @@ from transformers.models.paligemma.modeling_paligemma import PaliGemmaForConditi from lerobot.configs.policies import PreTrainedConfig from lerobot.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS -from lerobot.policies.pi05_openpi.configuration_pi05openpi import PI05OpenPIConfig +from lerobot.policies.pi05.configuration_pi05openpi import PI05OpenPIConfig from lerobot.policies.pretrained import PreTrainedPolicy, T diff --git a/src/lerobot/policies/pi0/processor_pi05openpi.py b/src/lerobot/policies/pi05/processor_pi05openpi.py similarity index 97% rename from src/lerobot/policies/pi0/processor_pi05openpi.py rename to src/lerobot/policies/pi05/processor_pi05openpi.py index 9f85db23c..5f993650f 100644 --- a/src/lerobot/policies/pi0/processor_pi05openpi.py +++ b/src/lerobot/policies/pi05/processor_pi05openpi.py @@ -6,8 +6,8 @@ import torch from lerobot.configs.types import PipelineFeatureType, PolicyFeature from lerobot.constants import OBS_STATE, POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME -from lerobot.policies.pi05_openpi.configuration_pi05openpi import PI05OpenPIConfig -from lerobot.policies.pi05_openpi.modeling_pi05openpi import pad_vector +from lerobot.policies.pi05.configuration_pi05openpi import PI05OpenPIConfig +from lerobot.policies.pi05.modeling_pi05openpi import pad_vector from lerobot.processor import ( AddBatchDimensionProcessorStep, DeviceProcessorStep, diff --git a/tests/policies/pi0_pi05/test_pi0_openpi.py b/tests/policies/pi0_pi05/test_pi0_openpi.py index 773cec56d..d87d06b5e 100644 --- a/tests/policies/pi0_pi05/test_pi0_openpi.py +++ b/tests/policies/pi0_pi05/test_pi0_openpi.py @@ -14,10 +14,10 @@ pytestmark = pytest.mark.skipif( ) from lerobot.policies.factory import make_policy_config # noqa: E402 -from lerobot.policies.pi0_openpi import ( # noqa: E402 - PI0OpenPIConfig, - PI0OpenPIPolicy, - make_pi0_openpi_pre_post_processors, # noqa: E402 +from lerobot.policies.pi0 import ( # noqa: E402 + PI0Config, + PI0Policy, + make_pi0_pre_post_processors, # noqa: E402 ) from lerobot.utils.random_utils import set_seed # noqa: E402 from tests.utils import require_cuda # noqa: E402 @@ -30,7 +30,7 @@ def test_policy_instantiation(): # Create config set_seed(42) - config = PI0OpenPIConfig(max_action_dim=7, max_state_dim=14, dtype="float32") + config = PI0Config(max_action_dim=7, max_state_dim=14, dtype="float32") # Set up input_features and output_features in the config from lerobot.configs.types import FeatureType, PolicyFeature @@ -70,10 +70,8 @@ def test_policy_instantiation(): } # Instantiate policy - policy = PI0OpenPIPolicy(config) - preprocessor, postprocessor = make_pi0_openpi_pre_post_processors( - config=config, dataset_stats=dataset_stats - ) + policy = PI0Policy(config) + preprocessor, postprocessor = make_pi0_pre_post_processors(config=config, dataset_stats=dataset_stats) # Test forward pass with dummy data batch_size = 1 device = config.device diff --git a/tests/policies/pi0_pi05/test_pi0_original_vs_lerobot.py b/tests/policies/pi0_pi05/test_pi0_original_vs_lerobot.py index 7b0f9712d..417148bb1 100644 --- a/tests/policies/pi0_pi05/test_pi0_original_vs_lerobot.py +++ b/tests/policies/pi0_pi05/test_pi0_original_vs_lerobot.py @@ -23,8 +23,8 @@ from openpi.models_pytorch import preprocessing_pytorch as openpi_preprocessing from openpi.models_pytorch.pi0_pytorch import PI0Pytorch # noqa: E402 from transformers import AutoTokenizer # noqa: E402 -from lerobot.policies.pi0_openpi import PI0OpenPIConfig, PI0OpenPIPolicy # noqa: E402 -from lerobot.policies.pi0_openpi.processor_pi0_openpi import make_pi0_openpi_pre_post_processors # noqa: E402 +from lerobot.policies.pi0 import PI0Config, PI0Policy # noqa: E402 +from lerobot.policies.pi0.processor_pi0_openpi import make_pi0_pre_post_processors # noqa: E402 from lerobot.processor import PolicyAction, PolicyProcessorPipeline # noqa: E402 # TODO: ADDING DEFAULT IMAGES_FEATURES TO CONFIG @@ -73,24 +73,20 @@ class PI0BaseOriginalConfig: def instantiate_lerobot_pi0( from_pretrained: bool = False, ) -> tuple[ - PI0OpenPIPolicy, + PI0Policy, PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], PolicyProcessorPipeline[PolicyAction, PolicyAction], ]: if from_pretrained: # Load the policy first - policy = PI0OpenPIPolicy.from_pretrained( - pretrained_name_or_path="pepijn223/pi0_base_fp32", strict=True - ) + policy = PI0Policy.from_pretrained(pretrained_name_or_path="pepijn223/pi0_base_fp32", strict=True) else: - config = PI0OpenPIConfig( - max_action_dim=DUMMY_ACTION_DIM, max_state_dim=DUMMY_STATE_DIM, dtype="float32" - ) - policy = PI0OpenPIPolicy(config) + config = PI0Config(max_action_dim=DUMMY_ACTION_DIM, max_state_dim=DUMMY_STATE_DIM, dtype="float32") + policy = PI0Policy(config) policy.to(DEVICE) policy.config.device = DEVICE - preprocessor, postprocessor = make_pi0_openpi_pre_post_processors( + preprocessor, postprocessor = make_pi0_pre_post_processors( config=policy.config, dataset_stats=DUMMY_DATASET_STATS ) return (policy, preprocessor, postprocessor)