|
|
--- |
|
|
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() |
|
|
``` |
