Skip to content

Commit 2a91b56

Browse files
STW-HSTW-H
committed
251208V0
1 parent 257db5f commit 2a91b56

11 files changed

Lines changed: 158 additions & 13 deletions

File tree

‎download.ipynb‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": null,
6+
"id": "9edfdf7f",
7+
"metadata": {},
8+
"outputs": [],
9+
"source": [
10+
"from lerobot.datasets.lerobot_dataset import LeRobotDataset\n",
11+
"\n",
12+
"ds = LeRobotDataset(repo_id=\"lerobot/aloha_sim_insertion_human\")"
13+
]
14+
}
15+
],
16+
"metadata": {
17+
"kernelspec": {
18+
"display_name": "openpi",
19+
"language": "python",
20+
"name": "python3"
21+
},
22+
"language_info": {
23+
"codemirror_mode": {
24+
"name": "ipython",
25+
"version": 3
26+
},
27+
"file_extension": ".py",
28+
"mimetype": "text/x-python",
29+
"name": "python",
30+
"nbconvert_exporter": "python",
31+
"pygments_lexer": "ipython3",
32+
"version": "3.11.14"
33+
}
34+
},
35+
"nbformat": 4,
36+
"nbformat_minor": 5
37+
}

‎examples/aloha_real/convert_aloha_data_to_lerobot.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,9 @@
1010
from typing import Literal
1111

1212
import h5py
13-
from lerobot.common.datasets.lerobot_dataset import LEROBOT_HOME
14-
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
15-
from lerobot.common.datasets.push_dataset_to_hub._download_raw import download_raw
13+
from lerobot.datasets.lerobot_dataset import LEROBOT_HOME
14+
from lerobot.datasets.lerobot_dataset import LeRobotDataset
15+
from lerobot.datasets.push_dataset_to_hub._download_raw import download_raw
1616
import numpy as np
1717
import torch
1818
import tqdm

‎examples/aloha_sim/main.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
class Args:
1616
out_dir: pathlib.Path = pathlib.Path("data/aloha_sim/videos")
1717

18-
task: str = "gym_aloha/AlohaTransferCube-v0"
18+
task: str = "gym_aloha/AlohaInsertion-v0"
1919
seed: int = 0
2020

2121
action_horizon: int = 10

‎examples/droid/convert_droid_data_to_lerobot.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@
1919

2020
import cv2
2121
import h5py
22-
from lerobot.common.datasets.lerobot_dataset import HF_LEROBOT_HOME
23-
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
22+
from lerobot.datasets.lerobot_dataset import HF_LEROBOT_HOME
23+
from lerobot.datasets.lerobot_dataset import LeRobotDataset
2424
import numpy as np
2525
from PIL import Image
2626
from tqdm import tqdm

‎examples/libero/convert_libero_data_to_lerobot.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,8 @@
2020

2121
import shutil
2222

23-
from lerobot.common.datasets.lerobot_dataset import HF_LEROBOT_HOME
24-
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
23+
from lerobot.datasets.lerobot_dataset import HF_LEROBOT_HOME
24+
from lerobot.datasets.lerobot_dataset import LeRobotDataset
2525
import tensorflow_datasets as tfds
2626
import tyro
2727

‎pyproject.toml‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ dependencies = [
1717
"imageio>=2.36.1",
1818
"jax[cuda12]==0.5.3",
1919
"jaxtyping==0.2.36",
20-
"lerobot",
20+
# "lerobot",
2121
"ml_collections==1.0.0",
2222
"numpy>=1.22.4,<2.0.0",
2323
"numpydantic>=1.6.6",
@@ -26,7 +26,7 @@ dependencies = [
2626
"orbax-checkpoint==0.11.13",
2727
"pillow>=11.0.0",
2828
"sentencepiece>=0.2.0",
29-
"torch==2.7.1",
29+
"torch==2.8.0+cu128",
3030
"tqdm-loggable>=0.2",
3131
"typing-extensions>=4.12.2",
3232
"tyro>=0.9.5",
@@ -61,10 +61,11 @@ rlds = [
6161

6262
[tool.uv]
6363
override-dependencies = ["ml-dtypes==0.4.1", "tensorstore==0.1.74"]
64+
extra-index-url = ["https://download.pytorch.org/whl/cu128"]
6465

6566
[tool.uv.sources]
6667
openpi-client = { workspace = true }
67-
lerobot = { git = "https://github.com/huggingface/lerobot", rev = "0cf864870cf29f4738d3ade893e6fd13fbd7cdb5" }
68+
# lerobot = { git = "https://github.com/huggingface/lerobot", rev = "0cf864870cf29f4738d3ade893e6fd13fbd7cdb5" }
6869
dlimp = { git = "https://github.com/kvablack/dlimp", rev = "ad72ce3a9b414db2185bc0b38461d4101a65477a" }
6970

7071
[tool.uv.workspace]

‎run.sh‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
#!/bin/bash
2+
3+
# Generate timestamp in YYYYMMDDHHMMSS format
4+
TIMESTAMP=$(date +"%Y%m%d%H%M%S")
5+
6+
# Training script for pi05 aloha simulation with automatic timestamp
7+
# export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9
8+
python scripts/train.py pi05_aloha_sim_insertion_human --exp-name=epri_aloha_sim_insertion_pi05${TIMESTAMP}
9+
# python scripts/train.py pi05_aloha_sim_transfer_cube_human --exp-name=epri_aloha_sim_insertion_pi05${TIMESTAMP} --resume

‎src/openpi/training/checkpoints.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,8 @@ def initialize_checkpoint_dir(
4848
max_to_keep=1,
4949
keep_period=keep_period,
5050
create=False,
51-
async_options=ocp.AsyncOptions(timeout_secs=7200),
51+
enable_async_checkpointing=False,
52+
# async_options=ocp.AsyncOptions(timeout_secs=7200),
5253
),
5354
)
5455

‎src/openpi/training/config.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -922,6 +922,46 @@ def __post_init__(self) -> None:
922922
num_train_steps=20_000,
923923
),
924924
#
925+
# ALOHA Sim pi05 configs. This config is used to demonstrate how to train on a simple simulated environment.
926+
#
927+
TrainConfig(
928+
name="pi05_aloha_sim_insertion_human",
929+
model=pi0_config.Pi0Config(pi05=True),
930+
data=LeRobotAlohaDataConfig(
931+
repo_id="lerobot/aloha_sim_insertion_human",
932+
default_prompt="Insert the peg into the socket.",
933+
use_delta_joint_actions=True,
934+
),
935+
lr_schedule=_optimizer.CosineDecaySchedule(
936+
warmup_steps=1_000,
937+
peak_lr=1e-5,
938+
decay_steps=50_000,
939+
decay_lr=1e-6,
940+
),
941+
942+
weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
943+
num_train_steps=50_000,
944+
945+
batch_size=32,
946+
num_workers=4,
947+
),
948+
TrainConfig(
949+
name="pi05_aloha_sim_transfer_cube_human",
950+
model=pi0_config.Pi0Config(pi05=True),
951+
data=LeRobotAlohaDataConfig(
952+
repo_id="lerobot/aloha_sim_transfer_cube_human",
953+
default_prompt="Transfer cube.",
954+
use_delta_joint_actions=False,
955+
),
956+
weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
957+
num_train_steps=10_000,
958+
959+
batch_size=12,
960+
num_workers=4,
961+
962+
# video_backend="pyav",
963+
),
964+
#
925965
# Debugging configs.
926966
#
927967
TrainConfig(

‎src/openpi/training/data_loader.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77

88
import jax
99
import jax.numpy as jnp
10-
import lerobot.common.datasets.lerobot_dataset as lerobot_dataset
10+
import lerobot.datasets.lerobot_dataset as lerobot_dataset
1111
import numpy as np
1212
import torch
1313

@@ -143,6 +143,7 @@ def create_torch_dataset(
143143
delta_timestamps={
144144
key: [t / dataset_meta.fps for t in range(action_horizon)] for key in data_config.action_sequence_keys
145145
},
146+
video_backend='pyav',
146147
)
147148

148149
if data_config.prompt_from_task:

0 commit comments

Comments
 (0)