debug_orbax.py

import sys
import pathlib
try:
    import orbax.checkpoint as ocp
except ImportError:
    print("Could not import orbax.checkpoint")
    sys.exit(1)

checkpoint_path = pathlib.Path("pi0_fast_base_params").resolve()

try:
    with ocp.PyTreeCheckpointer() as ckptr:
        metadata = ckptr.metadata(checkpoint_path)
        print(f"Meta type: {type(metadata)}")
        if hasattr(metadata, 'item_metadata'):
             print("Has item_metadata")
             # Try to print keys if it is a dict-like
             try:
                 print(f"item_metadata keys: {list(metadata.item_metadata.keys())}")
             except:
                 print(f"item_metadata type: {type(metadata.item_metadata)}")
        else:
             print("No item_metadata")
             
except Exception as e:
    print(f"Error: {e}")