File size: 1,095 Bytes
499bbbc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
---
license: mit
base_model:
- google/paligemma-3b-pt-224
tags:
- lerobot
- torch
- pi0
datasets:
- IPEC-COMMUNITY/bridge_orig_lerobot
---
download the model
```bash
huggingface-cli download --resume-download --local-dir-use-symlinks False ${model} --local-dir $(basename ${model})
```
install the [lerobot](https://github.com/huggingface/lerobot) package first
```python
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()
``` |