Skip to content

Restore z-score normalization for pi0_fast_libero configs #971

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
0xadvait wants to merge 1 commit into
base: main
Choose a base branch
from
+13 −1
Open

Restore z-score normalization for pi0_fast_libero configs #971

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion src/openpi/training/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,10 @@ class DataConfigFactory(abc.ABC):
assets: AssetsConfig = dataclasses.field(default_factory=AssetsConfig)
# Base config that will be updated by the factory.
base_config: tyro.conf.Suppress[DataConfig | None] = None
# If set, overrides the model-type-based default for quantile normalization.
# Released pi0-FAST LIBERO checkpoints were trained with z-score normalization,
# so their configs pin this to False for checkpoint compatibility.
use_quantile_norm_override: bool | None = None

@abc.abstractmethod
def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig) -> DataConfig:
Expand All @@ -184,7 +188,9 @@ def create_base_config(self, assets_dirs: pathlib.Path, model_config: _model.Bas
repo_id=repo_id,
asset_id=asset_id,
norm_stats=self._load_norm_stats(epath.Path(self.assets.assets_dir or assets_dirs), asset_id),
use_quantile_norm=model_config.model_type != ModelType.PI0,
use_quantile_norm=self.use_quantile_norm_override
if self.use_quantile_norm_override is not None
else model_config.model_type != ModelType.PI0,
)

def _load_norm_stats(self, assets_dir: epath.Path, asset_id: str | None) -> dict[str, _transforms.NormStats] | None:
Expand Down Expand Up @@ -713,6 +719,9 @@ def __post_init__(self) -> None:
repo_id="physical-intelligence/libero",
base_config=DataConfig(prompt_from_task=True),
extra_delta_transform=True,
# The released pi0_fast_libero checkpoint predates the quantile-norm default for non-PI0
# models and was trained with z-score normalization, so keep normalization z-score.
use_quantile_norm_override=False,
),
# Note that we load the pi0-FAST base model checkpoint here.
weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_fast_base/params"),
Expand All @@ -729,6 +738,9 @@ def __post_init__(self) -> None:
repo_id="physical-intelligence/libero",
base_config=DataConfig(prompt_from_task=True),
extra_delta_transform=True,
# The released pi0_fast_libero checkpoint predates the quantile-norm default for non-PI0
# models and was trained with z-score normalization, so keep normalization z-score.
use_quantile_norm_override=False,
),
weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_fast_base/params"),
num_train_steps=30_000,
Expand Down