Skip to content

Commit 263f85e

Browse files
juelgjuelg
committed
tmp
1 parent 8d95024 commit 263f85e

1 file changed

Lines changed: 261 additions & 0 deletions

File tree

‎python/rcs/lerobot_joint_converter.py‎

Lines changed: 261 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,261 @@
1+
from __future__ import annotations
2+
3+
from pathlib import Path
4+
5+
import duckdb
6+
import numpy as np
7+
import pandas as pd
8+
import pyarrow as pa
9+
import torch
10+
from torchvision.io import decode_jpeg
11+
from torchvision.transforms import v2
12+
13+
from lerobot.datasets.lerobot_dataset import LeRobotDataset
14+
15+
16+
DATASET_PATHS = [
17+
"data_grasp",
18+
]
19+
HF_DATA_DIR = "data_lerobot_joint_simple"
20+
REPO_ID = "rcs/grasp_joint_simple"
21+
ROBOT_TYPE = "fr3"
22+
FPS = 30
23+
ARM_KEY = "right"
24+
25+
CAMERAS = [
26+
("head", "head"),
27+
("image_left_wrist", "left_wrist"),
28+
("image_right_wrist", "right_wrist"),
29+
]
30+
IMAGE_SIZE = (256, 256)
31+
RESIZE = v2.Resize(IMAGE_SIZE)
32+
IMAGE_BATCH_SIZE = 32
33+
34+
35+
class JointDatasetConverter:
36+
def __init__(self, root: str | Path):
37+
self.root = Path(root)
38+
self.conn = duckdb.connect()
39+
self.source_sql = self._build_source_sql(DATASET_PATHS)
40+
41+
self.lrds = LeRobotDataset.create(
42+
repo_id=REPO_ID,
43+
robot_type=ROBOT_TYPE,
44+
root=self.root,
45+
fps=FPS,
46+
use_videos=False,
47+
features={
48+
"observation.images.head": {
49+
"dtype": "image",
50+
"shape": (*IMAGE_SIZE, 3),
51+
"names": ["height", "width", "channel"],
52+
},
53+
"observation.images.image_left_wrist": {
54+
"dtype": "image",
55+
"shape": (*IMAGE_SIZE, 3),
56+
"names": ["height", "width", "channel"],
57+
},
58+
"observation.images.image_right_wrist": {
59+
"dtype": "image",
60+
"shape": (*IMAGE_SIZE, 3),
61+
"names": ["height", "width", "channel"],
62+
},
63+
"observation.state": {
64+
"dtype": "float32",
65+
"shape": (8,),
66+
"names": [f"joint_{i}" for i in range(7)] + ["gripper"],
67+
},
68+
"action": {
69+
"dtype": "float32",
70+
"shape": (8,),
71+
"names": [f"joint_{i}" for i in range(7)] + ["gripper"],
72+
},
73+
},
74+
image_writer_threads=0,
75+
image_writer_processes=0,
76+
)
77+
78+
def _build_source_sql(self, dataset_paths: list[str | Path]) -> str:
79+
queries = []
80+
for path in dataset_paths:
81+
escaped = str(path).replace("'", "''")
82+
queries.append(f"SELECT * FROM read_parquet('{escaped}')")
83+
return " UNION ALL ".join(queries)
84+
85+
def generate_examples(self, success: bool = True, n: int = -1):
86+
uuids = self.conn.execute(
87+
f"SELECT DISTINCT uuid FROM ({self.source_sql}) AS src ORDER BY uuid"
88+
).fetchall()
89+
90+
for (episode_id,) in uuids:
91+
table = self._fetch_transition_table(episode_id)
92+
93+
converted = self.parse_episode(episode_id, table, success)
94+
if converted:
95+
n -= 1
96+
if n == 0:
97+
break
98+
99+
self.lrds.finalize()
100+
101+
def _fetch_transition_table(self, episode_id: str) -> pd.DataFrame:
102+
return self.conn.execute(
103+
f"""
104+
WITH ordered AS (
105+
SELECT
106+
uuid,
107+
step,
108+
success,
109+
instruction,
110+
obs.{ARM_KEY}.joints AS observation_joints,
111+
obs.{ARM_KEY}.gripper AS observation_gripper,
112+
action.{ARM_KEY}.gripper AS action_gripper,
113+
LEAD(obs.{ARM_KEY}.joints) OVER (PARTITION BY uuid ORDER BY step) AS next_joints,
114+
LEAD(obs.{ARM_KEY}.gripper) OVER (PARTITION BY uuid ORDER BY step) AS next_gripper
115+
FROM ({self.source_sql}) AS src
116+
WHERE uuid = ?
117+
)
118+
SELECT
119+
uuid,
120+
step,
121+
success,
122+
instruction,
123+
observation_joints,
124+
observation_gripper,
125+
action_gripper,
126+
next_joints,
127+
next_gripper
128+
FROM ordered
129+
WHERE next_joints IS NOT NULL
130+
AND NOT (
131+
observation_joints = next_joints
132+
AND observation_gripper = next_gripper
133+
)
134+
ORDER BY step
135+
""",
136+
[episode_id],
137+
).df()
138+
139+
def _fetch_episode_success(self, episode_id: str) -> bool:
140+
return bool(
141+
self.conn.execute(
142+
f"SELECT COALESCE(MAX(success), FALSE) FROM ({self.source_sql}) AS src WHERE uuid = ?",
143+
[episode_id],
144+
).fetchone()[0]
145+
)
146+
147+
def _image_query(self) -> str:
148+
return f"""
149+
WITH ordered AS (
150+
SELECT
151+
uuid,
152+
step,
153+
obs.{ARM_KEY}.joints AS observation_joints,
154+
obs.{ARM_KEY}.gripper AS observation_gripper,
155+
obs.frames.head.rgb.data AS image_head,
156+
obs.frames.left_wrist.rgb.data AS image_left_wrist,
157+
obs.frames.right_wrist.rgb.data AS image_right_wrist,
158+
LEAD(obs.{ARM_KEY}.joints) OVER (PARTITION BY uuid ORDER BY step) AS next_joints,
159+
LEAD(obs.{ARM_KEY}.gripper) OVER (PARTITION BY uuid ORDER BY step) AS next_gripper
160+
FROM ({self.source_sql}) AS src
161+
WHERE uuid = ?
162+
AND obs.frames.head.rgb.data IS NOT NULL
163+
AND obs.frames.left_wrist.rgb.data IS NOT NULL
164+
AND obs.frames.right_wrist.rgb.data IS NOT NULL
165+
)
166+
SELECT
167+
step,
168+
image_head,
169+
image_left_wrist,
170+
image_right_wrist
171+
FROM ordered
172+
WHERE next_joints IS NOT NULL
173+
AND NOT (
174+
observation_joints = next_joints
175+
AND observation_gripper = next_gripper
176+
)
177+
ORDER BY step
178+
"""
179+
180+
def parse_episode(self, episode_id: str, table: pd.DataFrame, success: bool):
181+
if len(table) == 0:
182+
return False
183+
184+
if success and not self._fetch_episode_success(episode_id):
185+
return False
186+
187+
df = table.reset_index(drop=True)
188+
rows_by_step = {int(row["step"]): row for _, row in df.iterrows()}
189+
step_order = [int(step) for step in df["step"].tolist()]
190+
frames_by_step: dict[int, dict[str, np.ndarray]] = {}
191+
192+
reader = self.conn.execute(self._image_query(), [episode_id]).fetch_record_batch(rows_per_batch=IMAGE_BATCH_SIZE)
193+
for batch in reader:
194+
self._decode_image_batch(batch, frames_by_step)
195+
196+
for step in step_order:
197+
curr = rows_by_step[step]
198+
images = frames_by_step[step]
199+
200+
state_vec = np.concatenate(
201+
[
202+
np.asarray(curr["observation_joints"], dtype=np.float32),
203+
np.asarray(curr["observation_gripper"], dtype=np.float32),
204+
]
205+
).astype(np.float32)
206+
207+
action_gripper = curr["action_gripper"]
208+
if action_gripper is None or action_gripper is pd.NA or (
209+
isinstance(action_gripper, float) and np.isnan(action_gripper)
210+
):
211+
action_gripper_vec = np.asarray(curr["next_gripper"], dtype=np.float32)
212+
else:
213+
action_gripper_vec = np.asarray(action_gripper, dtype=np.float32)
214+
215+
action_vec = np.concatenate(
216+
[
217+
np.asarray(curr["next_joints"], dtype=np.float32),
218+
action_gripper_vec,
219+
]
220+
).astype(np.float32)
221+
222+
self.lrds.add_frame(
223+
{
224+
"observation.images.head": images["head"],
225+
"observation.images.image_left_wrist": images["image_left_wrist"],
226+
"observation.images.image_right_wrist": images["image_right_wrist"],
227+
"observation.state": state_vec,
228+
"action": action_vec,
229+
"task": str(curr["instruction"]),
230+
}
231+
)
232+
233+
self.lrds.save_episode()
234+
return True
235+
236+
def _decode_and_resize_batch(self, image_bytes_list: list[bytes]) -> np.ndarray:
237+
image_tensors = [torch.frombuffer(bytearray(image_bytes), dtype=torch.uint8) for image_bytes in image_bytes_list]
238+
decoded = decode_jpeg(image_tensors)
239+
batch = torch.stack(decoded)
240+
resized = RESIZE(batch)
241+
return resized.permute(0, 2, 3, 1).cpu().numpy()
242+
243+
def _decode_image_batch(self, batch: pa.RecordBatch, frames_by_step: dict[int, dict[str, np.ndarray]]) -> None:
244+
batch_dict = batch.to_pydict()
245+
steps = [int(step) for step in batch_dict["step"]]
246+
decoded_images = {}
247+
for feature_name, column_name in CAMERAS:
248+
image_column = f"image_{column_name}" if column_name != "head" else "image_head"
249+
decoded_images[feature_name] = self._decode_and_resize_batch(batch_dict[image_column])
250+
251+
for idx, step in enumerate(steps):
252+
frames_by_step[step] = {
253+
"head": decoded_images["head"][idx],
254+
"image_left_wrist": decoded_images["image_left_wrist"][idx],
255+
"image_right_wrist": decoded_images["image_right_wrist"][idx],
256+
}
257+
258+
259+
if __name__ == "__main__":
260+
hf_ds = JointDatasetConverter(HF_DATA_DIR)
261+
hf_ds.generate_examples(success=True, n=-1)

0 commit comments

Comments
 (0)