Views
No views yet
25000pi05_cloth_folding_v4/home/prakhar-goel/data_farm/checkpoints/pi05/pi05_cloth_folding_v4/25000params/: Model parameters in JAX/Flax formattrain_state/: Training state including optimizer statesassets/: Normalization statistics and other training assets_CHECKPOINT_METADATA: Checkpoint metadata1from pathlib import Path
2import jax
3import orbax.checkpoint as ocp
4from etils import epath
5
6checkpoint_dir = Path("path/to/downloaded/checkpoint")
7
8# Load parameters
9restored = ocp.CheckpointManager(checkpoint_dir).restore(
10 checkpoint_dir,
11 items={
12 "params": ...,
13 "train_state": ...,
14 }
15)1from openpi.policies.policy_config import create_trained_policy
2
3# Load policy from checkpoint
4policy = create_trained_policy(
5 train_config=your_config,
6 checkpoint_dir=str(checkpoint_dir)
7)1{
2 "item_handlers": {
3 "assets": "openpi.training.checkpoints.CallbackHandler",
4 "params": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler",
5 "train_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"
6 },
7 "metrics": {},
8 "performance_metrics": {},
9 "init_timestamp_nsecs": 1762495032653103749,
10 "commit_timestamp_nsecs": 1762495109617216108,
11 "custom_metadata": {}
12}