Fix duplicated code from an example above

#3
by nikodembartnik HF Staff - opened
Files changed (1) hide show
  1. app/src/content/article.mdx +106 -43
app/src/content/article.mdx CHANGED
@@ -220,56 +220,119 @@ Record a Dataset
220
  [https://github.com/fracapuano/robot-learning-tutorial/blob/main/snippets/ch1/02_record_data.py](https://github.com/fracapuano/robot-learning-tutorial/blob/main/snippets/ch1/02_record_data.py)
221
 
222
  ```python
223
- import torch
 
 
 
 
224
  from lerobot.datasets.lerobot_dataset import LeRobotDataset
225
- from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
226
-
227
- delta_timestamps = {
228
- "observation.images.wrist_camera": [-0.2, -0.1, 0.0] # 0.2, and 0.1 seconds *before* each frame
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
229
  }
230
-
231
- # Optionally, use StreamingLeRobotDataset to avoid downloading the dataset
232
- dataset = LeRobotDataset(
233
- "lerobot/svla_so101_pickplace",
234
- delta_timestamps=delta_timestamps
235
  )
236
-
237
- # Streams frames from the Hugging Face Hub without loading into memory
238
- streaming_dataset = StreamingLeRobotDataset(
239
- "lerobot/svla_so101_pickplace",
240
- delta_timestamps=delta_timestamps
241
  )
242
 
243
- # Get the 100th frame in the dataset by
244
- sample = dataset[100]
245
- print(sample)
246
- # {
247
- # 'observation.state': tensor([...]),
248
- # 'action': tensor([...]),
249
- # 'observation.images.wrist_camera': tensor([3, C, H, W]), for delta timesteps
250
- # ...
251
- # }
252
-
253
- batch_size=16
254
- # wrap the dataset in a DataLoader to use process it batches for training purposes
255
- data_loader = torch.utils.data.DataLoader(
256
- dataset,
257
- batch_size=batch_size
 
 
258
  )
259
 
260
- # Iterate over the DataLoader in a training loop
261
- num_epochs = 1
262
- device = "cuda" if torch.cuda.is_available() else "cpu"
263
-
264
- for epoch in range(num_epochs):
265
- for batch in data_loader:
266
- # Move data to the appropriate device (e.g., GPU)
267
- observations = batch["observation.state"].to(device)
268
- actions = batch["action"].to(device)
269
- images = batch["observation.images.wrist_camera"].to(device)
270
-
271
- # Next, you can do amazing_model.forward(batch)
272
- ...
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
273
  ```
274
 
275
  </div>
 
220
  [https://github.com/fracapuano/robot-learning-tutorial/blob/main/snippets/ch1/02_record_data.py](https://github.com/fracapuano/robot-learning-tutorial/blob/main/snippets/ch1/02_record_data.py)
221
 
222
  ```python
223
+ """
224
+ You can also use the CLI to record data. To see the required arguments, run:
225
+ lerobot-record --help
226
+ """
227
+ from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
228
  from lerobot.datasets.lerobot_dataset import LeRobotDataset
229
+ from lerobot.datasets.utils import hw_to_dataset_features
230
+ from lerobot.robots.so100_follower import SO100Follower, SO100FollowerConfig
231
+ from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig
232
+ from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
233
+ from lerobot.utils.control_utils import init_keyboard_listener
234
+ from lerobot.utils.utils import log_say
235
+ from lerobot.utils.visualization_utils import init_rerun
236
+ from lerobot.scripts.lerobot_record import record_loop
237
+
238
+ NUM_EPISODES = 5
239
+ FPS = 30
240
+ EPISODE_TIME_SEC = 60
241
+ RESET_TIME_SEC = 10
242
+ TASK_DESCRIPTION = ... # provide a task description
243
+
244
+ HF_USER = ... # provide your Hugging Face username
245
+
246
+ follower_port = ... # find your ports running: lerobot-find-port
247
+ leader_port = ...
248
+ follower_id = ... # to load the calibration file
249
+ leader_id = ...
250
+
251
+ # Create the robot and teleoperator configurations
252
+ camera_config = {"front": OpenCVCameraConfig(
253
+ index_or_path=0, width=640, height=480, fps=FPS)
254
  }
255
+ robot_config = SO100FollowerConfig(
256
+ port=follower_port,
257
+ id=follower_id,
258
+ cameras=camera_config
 
259
  )
260
+ teleop_config = SO100LeaderConfig(
261
+ port=leader_port,
262
+ id=leader_id
 
 
263
  )
264
 
265
+ # Initialize the robot and teleoperator
266
+ robot = SO100Follower(robot_config)
267
+ teleop = SO100Leader(teleop_config)
268
+
269
+ # Configure the dataset features
270
+ action_features = hw_to_dataset_features(robot.action_features, "action")
271
+ obs_features = hw_to_dataset_features(robot.observation_features, "observation")
272
+ dataset_features = {**action_features, **obs_features}
273
+
274
+ # Create the dataset where to store the data
275
+ dataset = LeRobotDataset.create(
276
+ repo_id=f"{HF_USER}/robot-learning-tutorial-data",
277
+ fps=FPS,
278
+ features=dataset_features,
279
+ robot_type=robot.name,
280
+ use_videos=True,
281
+ image_writer_threads=4,
282
  )
283
 
284
+ # Initialize the keyboard listener and rerun visualization
285
+ _, events = init_keyboard_listener()
286
+ init_rerun(session_name="recording")
287
+
288
+ # Connect the robot and teleoperator
289
+ robot.connect()
290
+ teleop.connect()
291
+
292
+ episode_idx = 0
293
+ while episode_idx < NUM_EPISODES and not events["stop_recording"]:
294
+ log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
295
+
296
+ record_loop(
297
+ robot=robot,
298
+ events=events,
299
+ fps=FPS,
300
+ teleop=teleop,
301
+ dataset=dataset,
302
+ control_time_s=EPISODE_TIME_SEC,
303
+ single_task=TASK_DESCRIPTION,
304
+ display_data=True,
305
+ )
306
+
307
+ # Reset the environment if not stopping or re-recording
308
+ if (not events["stop_recording"]) and \
309
+ (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
310
+ log_say("Reset the environment")
311
+ record_loop(
312
+ robot=robot,
313
+ events=events,
314
+ fps=FPS,
315
+ teleop=teleop,
316
+ control_time_s=RESET_TIME_SEC,
317
+ single_task=TASK_DESCRIPTION,
318
+ display_data=True,
319
+ )
320
+
321
+ if events["rerecord_episode"]:
322
+ log_say("Re-recording episode")
323
+ events["rerecord_episode"] = False
324
+ events["exit_early"] = False
325
+ dataset.clear_episode_buffer()
326
+ continue
327
+
328
+ dataset.save_episode()
329
+ episode_idx += 1
330
+
331
+ # Clean up
332
+ log_say("Stop recording")
333
+ robot.disconnect()
334
+ teleop.disconnect()
335
+ dataset.push_to_hub()
336
  ```
337
 
338
  </div>