| """Load the released RoboTwin Pi05 checkpoint with portable asset paths.""" |
|
|
| from pathlib import Path |
|
|
| from openpi.models.pi0_config import Pi0Config |
| from openpi.policies.policy_config import create_trained_policy |
| from openpi.training.config import AssetsConfig, DataConfig, LeRobotAlohaDataConfig, TrainConfig |
|
|
| ASSET_ID = "RoboTwin-pi05_900-aloha_agilex-joint" |
| CONFIG_NAME = "pi05_base_aloha_full_sim_robotwin_900_defaults" |
|
|
|
|
| def make_config(checkpoint_dir): |
| checkpoint_dir = Path(checkpoint_dir).resolve() |
| return TrainConfig( |
| name=CONFIG_NAME, |
| model=Pi0Config(pi05=True), |
| data=LeRobotAlohaDataConfig( |
| repo_id=ASSET_ID, |
| assets=AssetsConfig( |
| assets_dir=str(checkpoint_dir / "assets"), |
| asset_id=ASSET_ID, |
| ), |
| base_config=DataConfig(prompt_from_task=True), |
| use_delta_joint_actions=True, |
| adapt_to_pi=False, |
| ), |
| ) |
|
|
|
|
| def load_policy(checkpoint_dir="14999"): |
| checkpoint_dir = Path(checkpoint_dir).resolve() |
| if not (checkpoint_dir / "params").is_dir(): |
| raise FileNotFoundError(f"Missing parameters: {checkpoint_dir / 'params'}") |
| stats = checkpoint_dir / "assets" / ASSET_ID / "norm_stats.json" |
| if not stats.is_file(): |
| raise FileNotFoundError(f"Missing normalization statistics: {stats}") |
| return create_trained_policy(make_config(checkpoint_dir), str(checkpoint_dir)) |
|
|
