lerobot-pi0-bridge / README.md
HaomingSong's picture
Create README.md (#1)
499bbbc verified
metadata
license: mit
base_model:
  - google/paligemma-3b-pt-224
tags:
  - lerobot
  - torch
  - pi0
datasets:
  - IPEC-COMMUNITY/bridge_orig_lerobot

download the model

huggingface-cli download --resume-download --local-dir-use-symlinks False ${model} --local-dir $(basename ${model})

install the lerobot package first

def auto_model_from_pretrained(path, **kwargs):
    import sys
    sys.path.append(path)  # noqa

    map_location = kwargs.pop("map_location", "cpu")
    from modeling_pi0 import PI0Policy
    return PI0Policy.from_pretrained(path, **kwargs).to(map_location)


policy = auto_model_from_pretrained(saved_model_path, map_location="cuda")

state = torch.rand(8)
image = np.array(Image.open("test/example.png"))
image = torch.from_numpy(image / 255).permute(2, 0, 1)

observation = {
    "observation.state": state.unsqueeze(0).to("cuda"),
    "observation.images.image_0": image.unsqueeze(0).to("cuda"),
    "task": ["put the object in the box"],
}

action_chunk = policy.select_action(observation)[0].cpu().numpy()