Skip to content

Commit a0cf9c7

Browse files
yhnsuyuanhaonanyuecideng
authored
Add LeRobot Data Handler for Dataset Export (#41)
Co-authored-by: yuanhaonan <yuanhaonan@dexforce.top> Co-authored-by: Yueci Deng <dengyueci@qq.com>
1 parent 83eea58 commit a0cf9c7

File tree

29 files changed

+1580
-92
lines changed

29 files changed

+1580
-92
lines changed

‎.github/workflows/main.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ jobs:
8080
export HF_ENDPOINT=https://hf-mirror.com
8181
pip uninstall pymeshlab -y
8282
pip install pymeshlab==2023.12.post3
83+
pip install numpy==1.26.4
8384
pytest tests
8485
8586
publish:

‎configs/gym/pour_water/gym_config.json‎

Lines changed: 33 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"id": "PourWater-v3",
3-
"max_episodes": 5,
3+
"max_episodes": 10,
44
"env": {
55
"events": {
66
"random_light": {
@@ -258,11 +258,38 @@
258258
}
259259
},
260260
"dataset": {
261-
"robot_meta": {
262-
"arm_dofs": 12,
263-
"control_freq": 25,
264-
"control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"],
265-
"min_len_steps": 5
261+
"lerobot": {
262+
"func": "LeRobotRecorder",
263+
"mode": "save",
264+
"params": {
265+
"robot_meta": {
266+
"robot_type": "CobotMagic",
267+
"arm_dofs": 12,
268+
"control_freq": 25,
269+
"control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"],
270+
"observation": {
271+
"vision": {
272+
"cam_high": ["mask"],
273+
"cam_right_wrist": ["mask"],
274+
"cam_left_wrist": ["mask"]
275+
},
276+
"states": ["qpos"],
277+
"exteroception": ["cam_high", "cam_right_wrist", "cam_left_wrist"]
278+
},
279+
"action": "qpos_with_eef_pose",
280+
"min_len_steps": 5
281+
},
282+
"instruction": {
283+
"lang": "Pour water from bottle to cup"
284+
},
285+
"extra": {
286+
"scene_type": "Commercial",
287+
"task_description": "Pour water",
288+
"data_type": "sim"
289+
},
290+
"use_videos": true,
291+
"export_success_only": false
292+
}
266293
}
267294
}
268295
},

‎configs/gym/pour_water/gym_config_simple.json‎

Lines changed: 326 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,326 @@
1+
{
2+
"id": "PourWater-v3",
3+
"max_episodes": 5,
4+
"env": {
5+
"events": {
6+
"random_light": {
7+
"func": "randomize_light",
8+
"mode": "interval",
9+
"interval_step": 10,
10+
"params": {
11+
"entity_cfg": {"uid": "light_1"},
12+
"position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]],
13+
"color_range": [[0.6, 0.6, 0.6], [1, 1, 1]],
14+
"intensity_range": [50.0, 100.0]
15+
}
16+
},
17+
"init_bottle_pose": {
18+
"func": "randomize_rigid_object_pose",
19+
"mode": "reset",
20+
"params": {
21+
"entity_cfg": {"uid": "bottle"},
22+
"position_range": [[-0.08, -0.12, 0.0], [0.08, 0.04, 0.0]],
23+
"relative_position": true
24+
}
25+
},
26+
"init_cup_pose": {
27+
"func": "randomize_rigid_object_pose",
28+
"mode": "reset",
29+
"params": {
30+
"entity_cfg": {"uid": "cup"},
31+
"position_range": [[-0.08, -0.04, 0.0], [0.08, 0.12, 0.0]],
32+
"relative_position": true
33+
}
34+
},
35+
"prepare_extra_attr": {
36+
"func": "prepare_extra_attr",
37+
"mode": "reset",
38+
"params": {
39+
"attrs": [
40+
{
41+
"name": "object_lengths",
42+
"mode": "callable",
43+
"entity_uids": "all_objects",
44+
"func_name": "compute_object_length",
45+
"func_kwargs": {
46+
"is_svd_frame": true,
47+
"sample_points": 5000
48+
}
49+
},
50+
{
51+
"name": "grasp_pose_object",
52+
"mode": "static",
53+
"entity_cfg": {
54+
"uid": "bottle"
55+
},
56+
"value": [[
57+
[0.32243, 0.03245, 0.94604, 0.025],
58+
[0.00706, -0.99947, 0.03188, -0.0 ],
59+
[0.94657, -0.0036 , -0.32249, 0.0 ],
60+
[0.0 , 0.0 , 0.0 , 1.0 ]
61+
]]
62+
},
63+
{
64+
"name": "left_arm_base_pose",
65+
"mode": "callable",
66+
"entity_cfg": {
67+
"uid": "CobotMagic"
68+
},
69+
"func_name": "get_link_pose",
70+
"func_kwargs": {
71+
"link_name": "left_arm_base",
72+
"to_matrix": true
73+
}
74+
},
75+
{
76+
"name": "right_arm_base_pose",
77+
"mode": "callable",
78+
"entity_cfg": {
79+
"uid": "CobotMagic"
80+
},
81+
"func_name": "get_link_pose",
82+
"func_kwargs": {
83+
"link_name": "right_arm_base",
84+
"to_matrix": true
85+
}
86+
}
87+
]
88+
}
89+
},
90+
"register_info_to_env": {
91+
"func": "register_info_to_env",
92+
"mode": "reset",
93+
"params": {
94+
"registry": [
95+
{
96+
"entity_cfg": {
97+
"uid": "bottle"
98+
},
99+
"pose_register_params": {
100+
"compute_relative": false,
101+
"compute_pose_object_to_arena": true,
102+
"to_matrix": true
103+
}
104+
},
105+
{
106+
"entity_cfg": {
107+
"uid": "cup"
108+
},
109+
"pose_register_params": {
110+
"compute_relative": false,
111+
"compute_pose_object_to_arena": true,
112+
"to_matrix": true
113+
}
114+
},
115+
{
116+
"entity_cfg": {
117+
"uid": "CobotMagic",
118+
"control_parts": ["left_arm"]
119+
},
120+
"attrs": ["left_arm_base_pose"],
121+
"pose_register_params": {
122+
"compute_relative": "cup",
123+
"compute_pose_object_to_arena": false,
124+
"to_matrix": true
125+
},
126+
"prefix": false
127+
},
128+
{
129+
"entity_cfg": {
130+
"uid": "CobotMagic",
131+
"control_parts": ["right_arm"]
132+
},
133+
"attrs": ["right_arm_base_pose"],
134+
"pose_register_params": {
135+
"compute_relative": "bottle",
136+
"compute_pose_object_to_arena": false,
137+
"to_matrix": true
138+
},
139+
"prefix": false
140+
}
141+
],
142+
"registration": "affordance_datas",
143+
"sim_update": true
144+
}
145+
},
146+
"random_material": {
147+
"func": "randomize_visual_material",
148+
"mode": "interval",
149+
"interval_step": 10,
150+
"params": {
151+
"entity_cfg": {"uid": "table"},
152+
"random_texture_prob": 0.0,
153+
"texture_path": "CocoBackground/coco",
154+
"base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]]
155+
}
156+
},
157+
"random_cup_material": {
158+
"func": "randomize_visual_material",
159+
"mode": "interval",
160+
"interval_step": 10,
161+
"params": {
162+
"entity_cfg": {"uid": "cup"},
163+
"random_texture_prob": 0.0,
164+
"texture_path": "CocoBackground/coco",
165+
"base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]]
166+
}
167+
},
168+
"random_bottle_material": {
169+
"func": "randomize_visual_material",
170+
"mode": "interval",
171+
"interval_step": 10,
172+
"params": {
173+
"entity_cfg": {"uid": "bottle"},
174+
"random_texture_prob": 0.0,
175+
"texture_path": "CocoBackground/coco",
176+
"base_color_range": [[0.2, 0.2, 0.2], [1.0, 1.0, 1.0]]
177+
}
178+
},
179+
"random_robot_init_eef_pose": {
180+
"func": "randomize_robot_eef_pose",
181+
"mode": "reset",
182+
"params": {
183+
"entity_cfg": {"uid": "CobotMagic", "control_parts": ["left_arm", "right_arm"]},
184+
"position_range": [[-0.01, -0.01, -0.01], [0.01, 0.01, 0]]
185+
}
186+
}
187+
},
188+
"observations": {
189+
"norm_robot_eef_joint": {
190+
"func": "normalize_robot_joint_data",
191+
"mode": "modify",
192+
"name": "robot/qpos",
193+
"params": {
194+
"joint_ids": [12, 13, 14, 15]
195+
}
196+
}
197+
},
198+
"dataset": {
199+
"lerobot": {
200+
"func": "LeRobotRecorder",
201+
"mode": "save",
202+
"params": {
203+
"robot_meta": {
204+
"robot_type": "CobotMagic",
205+
"arm_dofs": 12,
206+
"control_freq": 25,
207+
"control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"],
208+
"observation": {
209+
"vision": {
210+
"cam_high": ["mask"]
211+
},
212+
"states": ["qpos"],
213+
"exteroception": ["cam_high", "cam_right_wrist", "cam_left_wrist"]
214+
},
215+
"action": "qpos_with_eef_pose",
216+
"min_len_steps": 5
217+
},
218+
"instruction": {
219+
"lang": "Pour water from bottle to cup"
220+
},
221+
"extra": {
222+
"scene_type": "Commercial",
223+
"task_description": "Pour water",
224+
"data_type": "sim"
225+
},
226+
"use_videos": true,
227+
"export_success_only": false
228+
}
229+
}
230+
}
231+
},
232+
"robot": {
233+
"uid": "CobotMagic",
234+
"robot_type": "CobotMagic",
235+
"init_pos": [0.0, 0.0, 0.7775],
236+
"init_qpos": [-0.3,0.3,1.0,1.0,-1.2,-1.2,0.0,0.0,0.6,0.6,0.0,0.0,0.05,0.05,0.05,0.05]
237+
},
238+
"sensor": [
239+
{
240+
"sensor_type": "Camera",
241+
"uid": "cam_high",
242+
"width": 960,
243+
"height": 540,
244+
"intrinsics": [488.1665344238281, 488.1665344238281, 480, 270],
245+
"extrinsics": {
246+
"eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774],
247+
"target": [0.8586357573287919, 0, 0.5232553674540066],
248+
"up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347]
249+
}
250+
}
251+
],
252+
"light": {
253+
"direct": [
254+
{
255+
"uid": "light_1",
256+
"light_type": "point",
257+
"color": [1.0, 1.0, 1.0],
258+
"intensity": 50.0,
259+
"init_pos": [2, 0, 2],
260+
"radius": 10.0
261+
}
262+
]
263+
},
264+
"background": [
265+
{
266+
"uid": "table",
267+
"shape": {
268+
"shape_type": "Mesh",
269+
"fpath": "CircleTableSimple/circle_table_simple.ply",
270+
"compute_uv": true
271+
},
272+
"attrs" : {
273+
"mass": 10.0,
274+
"static_friction": 0.95,
275+
"dynamic_friction": 0.9,
276+
"restitution": 0.01
277+
},
278+
"body_scale": [1, 1, 1],
279+
"body_type": "kinematic",
280+
"init_pos": [0.725, 0.0, 0.825],
281+
"init_rot": [0, 90, 0]
282+
}
283+
],
284+
"rigid_object": [
285+
{
286+
"uid":"cup",
287+
"shape": {
288+
"shape_type": "Mesh",
289+
"fpath": "PaperCup/paper_cup.ply",
290+
"compute_uv": true
291+
},
292+
"attrs" : {
293+
"mass": 0.01,
294+
"contact_offset": 0.003,
295+
"rest_offset": 0.001,
296+
"restitution": 0.01,
297+
"max_depenetration_velocity": 1e1,
298+
"min_position_iters": 32,
299+
"min_velocity_iters":8
300+
},
301+
"init_pos": [0.75, 0.1, 0.9],
302+
"body_scale":[0.75, 0.75, 1.0],
303+
"max_convex_hull_num": 8
304+
},
305+
{
306+
"uid":"bottle",
307+
"shape": {
308+
"shape_type": "Mesh",
309+
"fpath": "ScannedBottle/kashijia_processed.ply",
310+
"compute_uv": true
311+
},
312+
"attrs" : {
313+
"mass": 0.01,
314+
"contact_offset": 0.003,
315+
"rest_offset": 0.001,
316+
"restitution": 0.01,
317+
"max_depenetration_velocity": 1e1,
318+
"min_position_iters": 32,
319+
"min_velocity_iters":8
320+
},
321+
"init_pos": [0.75, -0.1, 0.932],
322+
"body_scale":[1, 1, 1],
323+
"max_convex_hull_num": 8
324+
}
325+
]
326+
}

0 commit comments

Comments
 (0)