feat(autopolicy): draft v0.4

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>

README.md

@@ -23,15 +23,24 @@
 </div>
 
 <h2 align="center">
-    <p><a href="https://github.com/huggingface/lerobot/blob/main/examples/10_use_so100.md">New robot in town: SO-100</a></p>
+    <p><a href="https://github.com/huggingface/lerobot/blob/main/examples/10_use_so100.md">
+        Build Your Own SO-100 Robot!</a></p>
 </h2>
 
 <div align="center">
-    <img src="media/so100/leader_follower.webp?raw=true" alt="SO-100 leader and follower arms" title="SO-100 leader and follower arms" width="50%">
-    <p>We just added a new tutorial on how to build a more affordable robot, at the price of $110 per arm!</p>
-    <p>Teach it new skills by showing it a few moves with just a laptop.</p>
-    <p>Then watch your homemade robot act autonomously 🤯</p>
-    <p>Follow the link to the <a href="https://github.com/huggingface/lerobot/blob/main/examples/10_use_so100.md">full tutorial for SO-100</a>.</p>
+  <img src="media/so100/leader_follower.webp?raw=true" alt="SO-100 leader and follower arms" title="SO-100 leader and follower arms" width="50%">
+
+  <p><strong>Meet the SO-100 – Just $110 per arm!</strong></p>
+  <p>Train it in minutes with a few simple moves on your laptop.</p>
+  <p>Then sit back and watch your creation act autonomously! 🤯</p>
+
+  <p><a href="https://github.com/huggingface/lerobot/blob/main/examples/10_use_so100.md">
+      Get the full SO-100 tutorial here.</a></p>
+
+  <p>Want to take it to the next level? Make your SO-100 mobile by building LeKiwi!</p>
+  <p>Check out the <a href="https://github.com/huggingface/lerobot/blob/main/examples/11_use_lekiwi.md">LeKiwi tutorial</a> and bring your robot to life on wheels.</p>
+
+  <img src="media/lekiwi/kiwi.webp?raw=true" alt="LeKiwi mobile robot" title="LeKiwi mobile robot" width="50%">
 </div>
 
 <br/>
@@ -83,15 +92,20 @@ git clone https://github.com/huggingface/lerobot.git
 cd lerobot
 ```
 
-Create a virtual environment with Python 3.10 and activate it, e.g. with [`miniconda`](https://docs.anaconda.com/free/miniconda/index.html):
+Create a virtual environment with Python 3.10 and activate it using [`uv`](https://github.com/astral-sh/uv):
 ```bash
-conda create -y -n lerobot python=3.10
-conda activate lerobot
+# Install uv if you haven't already
+curl -LsSf https://astral.sh/uv/install.sh | sh
+
+# Create and activate virtual environment with Python 3.10
+uv venv .venv --python=3.10
+source .venv/bin/activate  # On Unix/macOS
+# .venv\Scripts\activate  # On Windows
 ```
 
 Install 🤗 LeRobot:
 ```bash
-pip install -e .
+uv pip install -e .
 ```
 
 > **NOTE:** Depending on your platform, If you encounter any build errors during this step

examples/10_use_so100.md

@@ -4,8 +4,8 @@
 
   - [A. Source the parts](#a-source-the-parts)
   - [B. Install LeRobot](#b-install-lerobot)
-  - [C. Configure the motors](#c-configure-the-motors)
-  - [D. Assemble the arms](#d-assemble-the-arms)
+  - [C. Configure the Motors](#c-configure-the-motors)
+  - [D. Step-by-Step Assembly Instructions](#d-step-by-step-assembly-instructions)
   - [E. Calibrate](#e-calibrate)
   - [F. Teleoperate](#f-teleoperate)
   - [G. Record a dataset](#g-record-a-dataset)
@@ -70,6 +70,7 @@ conda install -y -c conda-forge "opencv>=4.10.0"
 ```
 Great :hugs:! You are now done installing LeRobot and we can begin assembling the SO100 arms :robot:.
 Every time you now want to use LeRobot you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
+
 ## C. Configure the motors
 
 > [!NOTE]
@@ -98,22 +99,22 @@ Example output when identifying the leader arm's port (e.g., `/dev/tty.usbmodem5
 ```
 Finding all available ports for the MotorBus.
 ['/dev/tty.usbmodem575E0032081', '/dev/tty.usbmodem575E0031751']
-Remove the usb cable from your DynamixelMotorsBus and press Enter when done.
+Remove the usb cable from your MotorsBus and press Enter when done.
 
 [...Disconnect leader arm and press Enter...]
 
-The port of this DynamixelMotorsBus is /dev/tty.usbmodem575E0031751
+The port of this MotorsBus is /dev/tty.usbmodem575E0031751
 Reconnect the usb cable.
 ```
 Example output when identifying the follower arm's port (e.g., `/dev/tty.usbmodem575E0032081`, or possibly `/dev/ttyACM1` on Linux):
 ```
 Finding all available ports for the MotorBus.
 ['/dev/tty.usbmodem575E0032081', '/dev/tty.usbmodem575E0031751']
-Remove the usb cable from your DynamixelMotorsBus and press Enter when done.
+Remove the usb cable from your MotorsBus and press Enter when done.
 
 [...Disconnect follower arm and press Enter...]
 
-The port of this DynamixelMotorsBus is /dev/tty.usbmodem575E0032081
+The port of this MotorsBus is /dev/tty.usbmodem575E0032081
 Reconnect the usb cable.
 ```
 
@@ -221,19 +222,13 @@ Redo the process for all your motors until ID 6. Do the same for the 6 motors of
 
 Follow the video for removing gears. You need to remove the gear for the motors of the leader arm. As a result, you will only use the position encoding of the motor and reduce friction to more easily operate the leader arm.
 
-#### c. Add motor horn to all 12 motors
+## D. Step-by-Step Assembly Instructions
 
-<details>
-<summary><strong>Video adding motor horn</strong></summary>
+**Step 1: Clean Parts**
+- Remove all support material from the 3D-printed parts.
+---
 
-<video src="https://github.com/user-attachments/assets/ef3391a4-ad05-4100-b2bd-1699bf86c969"></video>
-
-</details>
-
-Follow the video for adding the motor horn. For SO-100, you need to align the holes on the motor horn to the motor spline to be approximately 1:30, 4:30, 7:30 and 10:30.
-Try to avoid rotating the motor while doing so to keep position 2048 set during configuration. It is especially tricky for the leader motors as it is more sensible without the gears, but it's ok if it's a bit rotated.
-
-## D. Assemble the arms
+### Additional Guidance
 
 <details>
 <summary><strong>Video assembling arms</strong></summary>
@@ -242,7 +237,211 @@ Try to avoid rotating the motor while doing so to keep position 2048 set during
 
 </details>
 
-Follow the video for assembling the arms. It is important to insert the cables into the motor that is being assembled before you assemble the motor into the arm! Inserting the cables beforehand is much easier than doing this afterward. The first arm should take a bit more than 1 hour to assemble, but once you get used to it, you can do it under 1 hour for the second arm.
+**Note:**
+This video provides visual guidance for assembling the arms, but it doesn't specify when or how to do the wiring. Inserting the cables beforehand is much easier than doing it afterward. The first arm may take a bit more than 1 hour to assemble, but once you get used to it, you can assemble the second arm in under 1 hour.
+
+---
+
+### First Motor
+
+**Step 2: Insert Wires**
+- Insert two wires into the first motor.
+
+  <img src="../media/tutorial/img1.jpg" style="height:300px;">
+
+**Step 3: Install in Base**
+- Place the first motor into the base.
+
+  <img src="../media/tutorial/img2.jpg" style="height:300px;">
+
+**Step 4: Secure Motor**
+- Fasten the motor with 4 screws. Two from the bottom and two from top.
+
+**Step 5: Attach Motor Holder**
+- Slide over the first motor holder and fasten it using two screws (one on each side).
+
+  <img src="../media/tutorial/img4.jpg" style="height:300px;">
+
+**Step 6: Attach Motor Horns**
+- Install both motor horns, securing the top horn with a screw. Try not to move the motor position when attaching the motor horn, especially for the leader arms, where we removed the gears.
+
+  <img src="../media/tutorial/img5.jpg" style="height:300px;">
+<details>
+  <summary><strong>Video adding motor horn</strong></summary>
+  <video src="https://github.com/user-attachments/assets/ef3391a4-ad05-4100-b2bd-1699bf86c969"></video>
+</details>
+
+**Step 7: Attach Shoulder Part**
+- Route one wire to the back of the robot and the other to the left or in photo towards you (see photo).
+- Attach the shoulder part.
+
+  <img src="../media/tutorial/img6.jpg" style="height:300px;">
+
+**Step 8: Secure Shoulder**
+- Tighten the shoulder part with 4 screws on top and 4 on the bottom
+*(access bottom holes by turning the shoulder).*
+
+---
+
+### Second Motor Assembly
+
+**Step 9: Install Motor 2**
+- Slide the second motor in from the top and link the wire from motor 1 to motor 2.
+
+  <img src="../media/tutorial/img8.jpg" style="height:300px;">
+
+**Step 10: Attach Shoulder Holder**
+- Add the shoulder motor holder.
+- Ensure the wire from motor 1 to motor 2 goes behind the holder while the other wire is routed upward (see photo).
+- This part can be tight to assemble, you can use a workbench like the image or a similar setup to push the part around the motor.
+
+  <div style="display: flex;">
+    <img src="../media/tutorial/img9.jpg" style="height:250px;">
+    <img src="../media/tutorial/img10.jpg" style="height:250px;">
+    <img src="../media/tutorial/img12.jpg" style="height:250px;">
+  </div>
+
+**Step 11: Secure Motor 2**
+- Fasten the second motor with 4 screws.
+
+**Step 12: Attach Motor Horn**
+- Attach both motor horns to motor 2, again use the horn screw.
+
+**Step 13: Attach Base**
+- Install the base attachment using 2 screws.
+
+  <img src="../media/tutorial/img11.jpg" style="height:300px;">
+
+**Step 14: Attach Upper Arm**
+- Attach the upper arm with 4 screws on each side.
+
+  <img src="../media/tutorial/img13.jpg" style="height:300px;">
+
+---
+
+### Third Motor Assembly
+
+**Step 15: Install Motor 3**
+- Route the motor cable from motor 2 through the cable holder to motor 3, then secure motor 3 with 4 screws.
+
+**Step 16: Attach Motor Horn**
+- Attach both motor horns to motor 3 and secure one again with a horn screw.
+
+  <img src="../media/tutorial/img14.jpg" style="height:300px;">
+
+**Step 17: Attach Forearm**
+- Connect the forearm to motor 3 using 4 screws on each side.
+
+  <img src="../media/tutorial/img15.jpg" style="height:300px;">
+
+---
+
+### Fourth Motor Assembly
+
+**Step 18: Install Motor 4**
+- Slide in motor 4, attach the cable from motor 3, and secure the cable in its holder with a screw.
+
+  <div style="display: flex;">
+    <img src="../media/tutorial/img16.jpg" style="height:300px;">
+    <img src="../media/tutorial/img19.jpg" style="height:300px;">
+  </div>
+
+**Step 19: Attach Motor Holder 4**
+- Install the fourth motor holder (a tight fit). Ensure one wire is routed upward and the wire from motor 3 is routed downward (see photo).
+
+  <img src="../media/tutorial/img17.jpg" style="height:300px;">
+
+**Step 20: Secure Motor 4 & Attach Horn**
+- Fasten motor 4 with 4 screws and attach its motor horns, use for one a horn screw.
+
+  <img src="../media/tutorial/img18.jpg" style="height:300px;">
+
+---
+
+### Wrist Assembly
+
+**Step 21: Install Motor 5**
+- Insert motor 5 into the wrist holder and secure it with 2 front screws.
+
+  <img src="../media/tutorial/img20.jpg" style="height:300px;">
+
+**Step 22: Attach Wrist**
+- Connect the wire from motor 4 to motor 5. And already insert the other wire for the gripper.
+- Secure the wrist to motor 4 using 4 screws on both sides.
+
+  <img src="../media/tutorial/img22.jpg" style="height:300px;">
+
+**Step 23: Attach Wrist Horn**
+- Install only one motor horn on the wrist motor and secure it with a horn screw.
+
+  <img src="../media/tutorial/img23.jpg" style="height:300px;">
+
+---
+
+### Follower Configuration
+
+**Step 24: Attach Gripper**
+- Attach the gripper to motor 5.
+
+  <img src="../media/tutorial/img24.jpg" style="height:300px;">
+
+**Step 25: Install Gripper Motor**
+- Insert the gripper motor, connect the motor wire from motor 5 to motor 6, and secure it with 3 screws on each side.
+
+  <img src="../media/tutorial/img25.jpg" style="height:300px;">
+
+**Step 26: Attach Gripper Horn & Claw**
+- Attach the motor horns and again use a horn screw.
+- Install the gripper claw and secure it with 4 screws on both sides.
+
+  <img src="../media/tutorial/img26.jpg" style="height:300px;">
+
+**Step 27: Mount Controller**
+- Attach the motor controller on the back.
+
+  <div style="display: flex;">
+    <img src="../media/tutorial/img27.jpg" style="height:300px;">
+    <img src="../media/tutorial/img28.jpg" style="height:300px;">
+  </div>
+
+*Assembly complete – proceed to Leader arm assembly.*
+
+---
+
+### Leader Configuration
+
+For the leader configuration, perform **Steps 1–23**. Make sure that you removed the motor gears from the motors.
+
+**Step 24: Attach Leader Holder**
+- Mount the leader holder onto the wrist and secure it with a screw.
+
+  <img src="../media/tutorial/img29.jpg" style="height:300px;">
+
+**Step 25: Attach Handle**
+- Attach the handle to motor 5 using 4 screws.
+
+  <img src="../media/tutorial/img30.jpg" style="height:300px;">
+
+**Step 26: Install Gripper Motor**
+- Insert the gripper motor, secure it with 3 screws on each side, attach a motor horn using a horn screw, and connect the motor wire.
+
+  <img src="../media/tutorial/img31.jpg" style="height:300px;">
+
+**Step 27: Attach Trigger**
+- Attach the follower trigger with 4 screws.
+
+  <img src="../media/tutorial/img32.jpg" style="height:300px;">
+
+**Step 28: Mount Controller**
+- Attach the motor controller on the back.
+
+  <div style="display: flex;">
+    <img src="../media/tutorial/img27.jpg" style="height:300px;">
+    <img src="../media/tutorial/img28.jpg" style="height:300px;">
+  </div>
+
+*Assembly complete – proceed to calibration.*
+
 
 ## E. Calibrate
 

examples/11_use_lekiwi.md

@@ -23,6 +23,9 @@ Follow this [README](https://github.com/SIGRobotics-UIUC/LeKiwi). It contains th
 
 Before assembling, you will first need to configure your motors. To this end, we provide a nice script, so let's first install LeRobot. After configuration, we will also guide you through assembly.
 
+### Wired version
+If you have the **wired** LeKiwi version you can skip the installation of the Raspberry Pi and setting up SSH. You can also run all commands directly on your PC for both the LeKiwi scripts and the leader arm scripts for teleoperating.
+
 ## B. Install software on Pi
 Now we have to setup the remote PC that will run on the LeKiwi Robot. This is normally a Raspberry Pi, but can be any PC that can run on 5V and has enough usb ports (2 or more) for the cameras and motor control board.
 
@@ -246,6 +249,110 @@ class LeKiwiRobotConfig(RobotConfig):
         }
     )
 
+    teleop_keys: dict[str, str] = field(
+        default_factory=lambda: {
+            # Movement
+            "forward": "w",
+            "backward": "s",
+            "left": "a",
+            "right": "d",
+            "rotate_left": "z",
+            "rotate_right": "x",
+            # Speed control
+            "speed_up": "r",
+            "speed_down": "f",
+            # quit teleop
+            "quit": "q",
+        }
+    )
+
+    mock: bool = False
+```
+
+## Wired version
+
+For the wired LeKiwi version your configured IP address should refer to your own laptop (127.0.0.1), because leader arm and LeKiwi are in this case connected to own laptop. Below and example configuration for this wired setup:
+```python
+@RobotConfig.register_subclass("lekiwi")
+@dataclass
+class LeKiwiRobotConfig(RobotConfig):
+    # `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
+    # Set this to a positive scalar to have the same value for all motors, or a list that is the same length as
+    # the number of motors in your follower arms.
+    max_relative_target: int | None = None
+
+    # Network Configuration
+    ip: str = "127.0.0.1"
+    port: int = 5555
+    video_port: int = 5556
+
+    cameras: dict[str, CameraConfig] = field(
+        default_factory=lambda: {
+            "front": OpenCVCameraConfig(
+                camera_index=0, fps=30, width=640, height=480, rotation=90
+            ),
+            "wrist": OpenCVCameraConfig(
+                camera_index=1, fps=30, width=640, height=480, rotation=180
+            ),
+        }
+    )
+
+    calibration_dir: str = ".cache/calibration/lekiwi"
+
+    leader_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "main": FeetechMotorsBusConfig(
+                port="/dev/tty.usbmodem585A0077581",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "sts3215"],
+                    "shoulder_lift": [2, "sts3215"],
+                    "elbow_flex": [3, "sts3215"],
+                    "wrist_flex": [4, "sts3215"],
+                    "wrist_roll": [5, "sts3215"],
+                    "gripper": [6, "sts3215"],
+                },
+            ),
+        }
+    )
+
+    follower_arms: dict[str, MotorsBusConfig] = field(
+        default_factory=lambda: {
+            "main": FeetechMotorsBusConfig(
+                port="/dev/tty.usbmodem58760431061",
+                motors={
+                    # name: (index, model)
+                    "shoulder_pan": [1, "sts3215"],
+                    "shoulder_lift": [2, "sts3215"],
+                    "elbow_flex": [3, "sts3215"],
+                    "wrist_flex": [4, "sts3215"],
+                    "wrist_roll": [5, "sts3215"],
+                    "gripper": [6, "sts3215"],
+                    "left_wheel": (7, "sts3215"),
+                    "back_wheel": (8, "sts3215"),
+                    "right_wheel": (9, "sts3215"),
+                },
+            ),
+        }
+    )
+
+    teleop_keys: dict[str, str] = field(
+        default_factory=lambda: {
+            # Movement
+            "forward": "w",
+            "backward": "s",
+            "left": "a",
+            "right": "d",
+            "rotate_left": "z",
+            "rotate_right": "x",
+            # Speed control
+            "speed_up": "r",
+            "speed_down": "f",
+            # quit teleop
+            "quit": "q",
+        }
+    )
+
     mock: bool = False
 ```
 
@@ -272,6 +379,9 @@ python lerobot/scripts/control_robot.py \
   --control.arms='["main_follower"]'
 ```
 
+### Wired version
+If you have the **wired** LeKiwi version please run all commands including this calibration command on your laptop.
+
 ### Calibrate leader arm
 Then to calibrate the leader arm (which is attached to the laptop/pc). You will need to move the leader arm to these positions sequentially:
 
@@ -326,6 +436,9 @@ You should see on your laptop something like this: ```[INFO] Connected to remote
 > [!TIP]
 >  If you use a different keyboard you can change the keys for each command in the [`LeKiwiRobotConfig`](../lerobot/common/robot_devices/robots/configs.py).
 
+### Wired version
+If you have the **wired** LeKiwi version please run all commands including both these teleoperation commands on your laptop.
+
 ## Troubleshoot communication
 
 If you are having trouble connecting to the Mobile SO100, follow these steps to diagnose and resolve the issue.
@@ -364,6 +477,13 @@ Make sure the configuration file on both your laptop/pc and the Raspberry Pi is
 # G. Record a dataset
 Once you're familiar with teleoperation, you can record your first dataset with LeKiwi.
 
+To start the program on LeKiwi, SSH into your Raspberry Pi, and run `conda activate lerobot` and this script:
+```bash
+python lerobot/scripts/control_robot.py \
+  --robot.type=lekiwi \
+  --control.type=remote_robot
+```
+
 If you want to use the Hugging Face hub features for uploading your dataset and you haven't previously done it, make sure you've logged in using a write-access token, which can be generated from the [Hugging Face settings](https://huggingface.co/settings/tokens):
 ```bash
 huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
@@ -374,8 +494,7 @@ Store your Hugging Face repository name in a variable to run these commands:
 HF_USER=$(huggingface-cli whoami | head -n 1)
 echo $HF_USER
 ```
-
-Record 2 episodes and upload your dataset to the hub:
+On your laptop then run this command to record 2 episodes and upload your dataset to the hub:
 ```bash
 python lerobot/scripts/control_robot.py \
   --robot.type=lekiwi \
@@ -393,6 +512,9 @@ python lerobot/scripts/control_robot.py \
 
 Note: You can resume recording by adding `--control.resume=true`.
 
+### Wired version
+If you have the **wired** LeKiwi version please run all commands including both these record dataset commands on your laptop.
+
 # H. Visualize a dataset
 
 If you uploaded your dataset to the hub with `--control.push_to_hub=true`, you can [visualize your dataset online](https://huggingface.co/spaces/lerobot/visualize_dataset) by copy pasting your repo id given by:

examples/3_train_policy.py

@@ -85,7 +85,7 @@ def main():
     done = False
     while not done:
         for batch in dataloader:
-            batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
+            batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
             loss, _ = policy.forward(batch)
             loss.backward()
             optimizer.step()

examples/port_datasets/droid_rlds/README.md

@@ -1,144 +0,0 @@
-# Port DROID 1.0.1 dataset to LeRobotDataset
-
-## Download
-
-TODO
-
-It will take 2 TB in your local disk.
-
-## Port on a single computer
-
-First, install tensorflow dataset utilities to read from raw files:
-```bash
-pip install tensorflow
-pip install tensorflow_datasets
-```
-
-Then run this script to start porting the dataset:
-```bash
-python examples/port_datasets/droid_rlds/port_droid.py \
-    --raw-dir /your/data/droid/1.0.1 \
-    --repo-id your_id/droid_1.0.1 \
-    --push-to-hub
-```
-
-It will take 400GB in your local disk.
-
-As usual, your LeRobotDataset will be stored in your huggingface/lerobot cache folder.
-
-WARNING: it will take 7 days for porting the dataset locally and 3 days to upload, so we will need to parallelize over multiple nodes on a slurm cluster.
-
-NOTE: For development, run this script to start porting a shard:
-```bash
-python examples/port_datasets/droid_rlds/port.py \
-    --raw-dir /your/data/droid/1.0.1 \
-    --repo-id your_id/droid_1.0.1 \
-    --num-shards 2048 \
-    --shard-index 0
-```
-
-## Port over SLURM
-
-Install slurm utilities from Hugging Face:
-```bash
-pip install datatrove
-```
-
-
-### 1. Port one shard per job
-
-Run this script to start porting shards of the dataset:
-```bash
-python examples/port_datasets/droid_rlds/slurm_port_shards.py \
-    --raw-dir /your/data/droid/1.0.1 \
-    --repo-id your_id/droid_1.0.1 \
-    --logs-dir /your/logs \
-    --job-name port_droid \
-    --partition your_partition \
-    --workers 2048 \
-    --cpus-per-task 8 \
-    --mem-per-cpu 1950M
-```
-
-**Note on how to set your command line arguments**
-
-Regarding `--partition`, find yours by running:
-```bash
-info --format="%R"`
-```
-and select the CPU partition if you have one. No GPU needed.
-
-Regarding `--workers`, it is the number of slurm jobs you will launch in parallel. 2048 is the maximum number, since there is 2048 shards in Droid. This big number will certainly max-out your cluster.
-
-Regarding `--cpus-per-task` and `--mem-per-cpu`, by default it will use ~16GB of RAM (8*1950M) which is recommended to load the raw frames and 8 CPUs which can be useful to parallelize the encoding of the frames.
-
-Find the number of CPUs and Memory of the nodes of your partition by running:
-```bash
-sinfo -N -p your_partition -h -o "%N cpus=%c mem=%m"
-```
-
-**Useful commands to check progress and debug**
-
-Check if your jobs are running:
-```bash
-squeue -u $USER`
-```
-
-You should see a list with job indices like `15125385_155` where `15125385` is the index of the run and `155` is the worker index. The output/print of this worker is written in real time in `/your/logs/job_name/slurm_jobs/15125385_155.out`. For instance, you can inspect the content of this file by running `less /your/logs/job_name/slurm_jobs/15125385_155.out`.
-
-Check the progression of your jobs by running:
-```bash
-jobs_status /your/logs
-```
-
-If it's not 100% and no more slurm job is running, it means that some of them failed. Inspect the logs by running:
-```bash
-failed_logs /your/logs/job_name
-```
-
-If there is an issue in the code, you can fix it in debug mode with `--slurm 0` which allows to set breakpoint:
-```bash
-python examples/port_datasets/droid_rlds/slurm_port_shards.py --slurm 0 ...
-```
-
-And you can relaunch the same command, which will skip the completed jobs:
-```bash
-python examples/port_datasets/droid_rlds/slurm_port_shards.py --slurm 1 ...
-```
-
-Once all jobs are completed, you will have one dataset per shard (e.g. `droid_1.0.1_world_2048_rank_1594`) saved on disk in your `/lerobot/home/dir/your_id` directory. You can find your `/lerobot/home/dir` by running:
-```bash
-python -c "from lerobot.common.constants import HF_LEROBOT_HOME;print(HF_LEROBOT_HOME)"
-```
-
-
-### 2. Aggregate all shards
-
-Run this script to start aggregation:
-```bash
-python examples/port_datasets/droid_rlds/slurm_aggregate_shards.py \
-    --repo-id your_id/droid_1.0.1 \
-    --logs-dir /your/logs \
-    --job-name aggr_droid \
-    --partition your_partition \
-    --workers 2048 \
-    --cpus-per-task 8 \
-    --mem-per-cpu 1950M
-```
-
-Once all jobs are completed, you will have one dataset your `/lerobot/home/dir/your_id/droid_1.0.1` directory.
-
-
-### 3. Upload dataset
-
-Run this script to start uploading:
-```bash
-python examples/port_datasets/droid_rlds/slurm_upload.py \
-    --repo-id your_id/droid_1.0.1 \
-    --logs-dir /your/logs \
-    --job-name upload_droid \
-    --partition your_partition \
-    --workers 50 \
-    --cpus-per-task 4 \
-    --mem-per-cpu 1950M
-```

examples/port_datasets/droid_rlds/port_droid.py

@@ -1,411 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import argparse
-import logging
-import time
-from pathlib import Path
-
-import numpy as np
-import tensorflow_datasets as tfds
-
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.utils.utils import get_elapsed_time_in_days_hours_minutes_seconds
-
-DROID_SHARDS = 2048
-DROID_FPS = 15
-DROID_ROBOT_TYPE = "Franka"
-
-# Dataset schema slightly adapted from: https://droid-dataset.github.io/droid/the-droid-dataset.html#-dataset-schema
-DROID_FEATURES = {
-    # true on first step of the episode
-    "is_first": {
-        "dtype": "bool",
-        "shape": (1,),
-        "names": None,
-    },
-    # true on last step of the episode
-    "is_last": {
-        "dtype": "bool",
-        "shape": (1,),
-        "names": None,
-    },
-    # true on last step of the episode if it is a terminal step, True for demos
-    "is_terminal": {
-        "dtype": "bool",
-        "shape": (1,),
-        "names": None,
-    },
-    # language_instruction is also stored as "task" to follow LeRobot standard
-    "language_instruction": {
-        "dtype": "string",
-        "shape": (1,),
-        "names": None,
-    },
-    "language_instruction_2": {
-        "dtype": "string",
-        "shape": (1,),
-        "names": None,
-    },
-    "language_instruction_3": {
-        "dtype": "string",
-        "shape": (1,),
-        "names": None,
-    },
-    "observation.state.gripper_position": {
-        "dtype": "float32",
-        "shape": (1,),
-        "names": {
-            "axes": ["gripper"],
-        },
-    },
-    "observation.state.cartesian_position": {
-        "dtype": "float32",
-        "shape": (6,),
-        "names": {
-            "axes": ["x", "y", "z", "roll", "pitch", "yaw"],
-        },
-    },
-    "observation.state.joint_position": {
-        "dtype": "float32",
-        "shape": (7,),
-        "names": {
-            "axes": ["x", "y", "z", "roll", "pitch", "yaw"],
-        },
-    },
-    # Add this new feature to follow LeRobot standard of using joint position + gripper
-    "observation.state": {
-        "dtype": "float32",
-        "shape": (8,),
-        "names": {
-            "axes": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6", "gripper"],
-        },
-    },
-    # Initially called wrist_image_left
-    "observation.images.wrist_left": {
-        "dtype": "video",
-        "shape": (180, 320, 3),
-        "names": [
-            "height",
-            "width",
-            "channels",
-        ],
-    },
-    # Initially called exterior_image_1_left
-    "observation.images.exterior_1_left": {
-        "dtype": "video",
-        "shape": (180, 320, 3),
-        "names": [
-            "height",
-            "width",
-            "channels",
-        ],
-    },
-    # Initially called exterior_image_2_left
-    "observation.images.exterior_2_left": {
-        "dtype": "video",
-        "shape": (180, 320, 3),
-        "names": [
-            "height",
-            "width",
-            "channels",
-        ],
-    },
-    "action.gripper_position": {
-        "dtype": "float32",
-        "shape": (1,),
-        "names": {
-            "axes": ["gripper"],
-        },
-    },
-    "action.gripper_velocity": {
-        "dtype": "float32",
-        "shape": (1,),
-        "names": {
-            "axes": ["gripper"],
-        },
-    },
-    "action.cartesian_position": {
-        "dtype": "float32",
-        "shape": (6,),
-        "names": {
-            "axes": ["x", "y", "z", "roll", "pitch", "yaw"],
-        },
-    },
-    "action.cartesian_velocity": {
-        "dtype": "float32",
-        "shape": (6,),
-        "names": {
-            "axes": ["x", "y", "z", "roll", "pitch", "yaw"],
-        },
-    },
-    "action.joint_position": {
-        "dtype": "float32",
-        "shape": (7,),
-        "names": {
-            "axes": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6"],
-        },
-    },
-    "action.joint_velocity": {
-        "dtype": "float32",
-        "shape": (7,),
-        "names": {
-            "axes": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6"],
-        },
-    },
-    # This feature was called "action" in RLDS dataset and consists of [6x joint velocities, 1x gripper position]
-    "action.original": {
-        "dtype": "float32",
-        "shape": (7,),
-        "names": {
-            "axes": ["x", "y", "z", "roll", "pitch", "yaw", "gripper"],
-        },
-    },
-    # Add this new feature to follow LeRobot standard of using joint position + gripper
-    "action": {
-        "dtype": "float32",
-        "shape": (8,),
-        "names": {
-            "axes": ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6", "gripper"],
-        },
-    },
-    "discount": {
-        "dtype": "float32",
-        "shape": (1,),
-        "names": None,
-    },
-    "reward": {
-        "dtype": "float32",
-        "shape": (1,),
-        "names": None,
-    },
-    # Meta data that are the same for all frames in the episode
-    "task_category": {
-        "dtype": "string",
-        "shape": (1,),
-        "names": None,
-    },
-    "building": {
-        "dtype": "string",
-        "shape": (1,),
-        "names": None,
-    },
-    "collector_id": {
-        "dtype": "string",
-        "shape": (1,),
-        "names": None,
-    },
-    "date": {
-        "dtype": "string",
-        "shape": (1,),
-        "names": None,
-    },
-    "camera_extrinsics.wrist_left": {
-        "dtype": "float32",
-        "shape": (6,),
-        "names": {
-            "axes": ["x", "y", "z", "roll", "pitch", "yaw"],
-        },
-    },
-    "camera_extrinsics.exterior_1_left": {
-        "dtype": "float32",
-        "shape": (6,),
-        "names": {
-            "axes": ["x", "y", "z", "roll", "pitch", "yaw"],
-        },
-    },
-    "camera_extrinsics.exterior_2_left": {
-        "dtype": "float32",
-        "shape": (6,),
-        "names": {
-            "axes": ["x", "y", "z", "roll", "pitch", "yaw"],
-        },
-    },
-    "is_episode_successful": {
-        "dtype": "bool",
-        "shape": (1,),
-        "names": None,
-    },
-}
-
-
-def is_episode_successful(tf_episode_metadata):
-    # Adapted from: https://github.com/droid-dataset/droid_policy_learning/blob/dd1020eb20d981f90b5ff07dc80d80d5c0cb108b/robomimic/utils/rlds_utils.py#L8
-    return "/success/" in tf_episode_metadata["file_path"].numpy().decode()
-
-
-def generate_lerobot_frames(tf_episode):
-    m = tf_episode["episode_metadata"]
-    frame_meta = {
-        "task_category": m["building"].numpy().decode(),
-        "building": m["building"].numpy().decode(),
-        "collector_id": m["collector_id"].numpy().decode(),
-        "date": m["date"].numpy().decode(),
-        "camera_extrinsics.wrist_left": m["extrinsics_wrist_cam"].numpy(),
-        "camera_extrinsics.exterior_1_left": m["extrinsics_exterior_cam_1"].numpy(),
-        "camera_extrinsics.exterior_2_left": m["extrinsics_exterior_cam_2"].numpy(),
-        "is_episode_successful": np.array([is_episode_successful(m)]),
-    }
-    for f in tf_episode["steps"]:
-        # Dataset schema slightly adapted from: https://droid-dataset.github.io/droid/the-droid-dataset.html#-dataset-schema
-        frame = {
-            "is_first": np.array([f["is_first"].numpy()]),
-            "is_last": np.array([f["is_last"].numpy()]),
-            "is_terminal": np.array([f["is_terminal"].numpy()]),
-            "language_instruction": f["language_instruction"].numpy().decode(),
-            "language_instruction_2": f["language_instruction_2"].numpy().decode(),
-            "language_instruction_3": f["language_instruction_3"].numpy().decode(),
-            "observation.state.gripper_position": f["observation"]["gripper_position"].numpy(),
-            "observation.state.cartesian_position": f["observation"]["cartesian_position"].numpy(),
-            "observation.state.joint_position": f["observation"]["joint_position"].numpy(),
-            "observation.images.wrist_left": f["observation"]["wrist_image_left"].numpy(),
-            "observation.images.exterior_1_left": f["observation"]["exterior_image_1_left"].numpy(),
-            "observation.images.exterior_2_left": f["observation"]["exterior_image_2_left"].numpy(),
-            "action.gripper_position": f["action_dict"]["gripper_position"].numpy(),
-            "action.gripper_velocity": f["action_dict"]["gripper_velocity"].numpy(),
-            "action.cartesian_position": f["action_dict"]["cartesian_position"].numpy(),
-            "action.cartesian_velocity": f["action_dict"]["cartesian_velocity"].numpy(),
-            "action.joint_position": f["action_dict"]["joint_position"].numpy(),
-            "action.joint_velocity": f["action_dict"]["joint_velocity"].numpy(),
-            "discount": np.array([f["discount"].numpy()]),
-            "reward": np.array([f["reward"].numpy()]),
-            "action.original": f["action"].numpy(),
-        }
-
-        # language_instruction is also stored as "task" to follow LeRobot standard
-        frame["task"] = frame["language_instruction"]
-
-        # Add this new feature to follow LeRobot standard of using joint position + gripper
-        frame["observation.state"] = np.concatenate(
-            [frame["observation.state.joint_position"], frame["observation.state.gripper_position"]]
-        )
-        frame["action"] = np.concatenate([frame["action.joint_position"], frame["action.gripper_position"]])
-
-        # Meta data that are the same for all frames in the episode
-        frame.update(frame_meta)
-
-        # Cast fp64 to fp32
-        for key in frame:
-            if isinstance(frame[key], np.ndarray) and frame[key].dtype == np.float64:
-                frame[key] = frame[key].astype(np.float32)
-
-        yield frame
-
-
-def port_droid(
-    raw_dir: Path,
-    repo_id: str,
-    push_to_hub: bool = False,
-    num_shards: int | None = None,
-    shard_index: int | None = None,
-):
-    dataset_name = raw_dir.parent.name
-    version = raw_dir.name
-    data_dir = raw_dir.parent.parent
-
-    builder = tfds.builder(f"{dataset_name}/{version}", data_dir=data_dir, version="")
-
-    if num_shards is not None:
-        tfds_num_shards = builder.info.splits["train"].num_shards
-        if tfds_num_shards != DROID_SHARDS:
-            raise ValueError(
-                f"Number of shards of Droid dataset is expected to be {DROID_SHARDS} but is {tfds_num_shards}."
-            )
-        if num_shards != tfds_num_shards:
-            raise ValueError(
-                f"We only shard over the fixed number of shards provided by tensorflow dataset ({tfds_num_shards}), but {num_shards} shards provided instead."
-            )
-        if shard_index >= tfds_num_shards:
-            raise ValueError(
-                f"Shard index is greater than the num of shards ({shard_index} >= {num_shards})."
-            )
-
-        raw_dataset = builder.as_dataset(split=f"train[{shard_index}shard]")
-    else:
-        raw_dataset = builder.as_dataset(split="train")
-
-    lerobot_dataset = LeRobotDataset.create(
-        repo_id=repo_id,
-        robot_type=DROID_ROBOT_TYPE,
-        fps=DROID_FPS,
-        features=DROID_FEATURES,
-    )
-
-    start_time = time.time()
-    num_episodes = raw_dataset.cardinality().numpy().item()
-    logging.info(f"Number of episodes {num_episodes}")
-
-    for episode_index, episode in enumerate(raw_dataset):
-        elapsed_time = time.time() - start_time
-        d, h, m, s = get_elapsed_time_in_days_hours_minutes_seconds(elapsed_time)
-
-        logging.info(
-            f"{episode_index} / {num_episodes} episodes processed (after {d} days, {h} hours, {m} minutes, {s:.3f} seconds)"
-        )
-
-        for frame in generate_lerobot_frames(episode):
-            lerobot_dataset.add_frame(frame)
-
-        lerobot_dataset.save_episode()
-        logging.info("Save_episode")
-
-    if push_to_hub:
-        lerobot_dataset.push_to_hub(
-            # Add openx tag, since it belongs to the openx collection of datasets
-            tags=["openx"],
-            private=False,
-        )
-
-
-def main():
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--raw-dir",
-        type=Path,
-        required=True,
-        help="Directory containing input raw datasets (e.g. `path/to/dataset` or `path/to/dataset/version).",
-    )
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        help="Repositery identifier on Hugging Face: a community or a user name `/` the name of the dataset, required when push-to-hub is True",
-    )
-    parser.add_argument(
-        "--push-to-hub",
-        action="store_true",
-        help="Upload to hub.",
-    )
-    parser.add_argument(
-        "--num-shards",
-        type=int,
-        default=None,
-        help="Number of shards. Can be either None to load the full dataset, or 2048 to load one of the 2048 tensorflow dataset files.",
-    )
-    parser.add_argument(
-        "--shard-index",
-        type=int,
-        default=None,
-        help="Index of the shard. Can be either None to load the full dataset, or in [0,2047] to load one of the 2048 tensorflow dataset files.",
-    )
-
-    args = parser.parse_args()
-
-    port_droid(**vars(args))
-
-
-if __name__ == "__main__":
-    main()

examples/port_datasets/droid_rlds/slurm_aggregate_shards.py

@@ -1,288 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import argparse
-import logging
-from pathlib import Path
-
-import tqdm
-from datatrove.executor import LocalPipelineExecutor
-from datatrove.executor.slurm import SlurmPipelineExecutor
-from datatrove.pipeline.base import PipelineStep
-
-from examples.port_datasets.droid_rlds.port_droid import DROID_SHARDS
-from lerobot.common.datasets.aggregate import validate_all_metadata
-from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
-from lerobot.common.datasets.utils import write_episode, write_episode_stats, write_info, write_task
-from lerobot.common.utils.utils import init_logging
-
-
-class AggregateDatasets(PipelineStep):
-    def __init__(
-        self,
-        repo_ids: list[str],
-        aggregated_repo_id: str,
-    ):
-        super().__init__()
-        self.repo_ids = repo_ids
-        self.aggr_repo_id = aggregated_repo_id
-
-        self.create_aggr_dataset()
-
-    def create_aggr_dataset(self):
-        init_logging()
-
-        logging.info("Start aggregate_datasets")
-
-        all_metadata = [LeRobotDatasetMetadata(repo_id) for repo_id in self.repo_ids]
-
-        fps, robot_type, features = validate_all_metadata(all_metadata)
-
-        # Create resulting dataset folder
-        aggr_meta = LeRobotDatasetMetadata.create(
-            repo_id=self.aggr_repo_id,
-            fps=fps,
-            robot_type=robot_type,
-            features=features,
-        )
-
-        logging.info("Find all tasks")
-        # find all tasks, deduplicate them, create new task indices for each dataset
-        # indexed by dataset index
-        datasets_task_index_to_aggr_task_index = {}
-        aggr_task_index = 0
-        for dataset_index, meta in enumerate(tqdm.tqdm(all_metadata, desc="Find all tasks")):
-            task_index_to_aggr_task_index = {}
-
-            for task_index, task in meta.tasks.items():
-                if task not in aggr_meta.task_to_task_index:
-                    # add the task to aggr tasks mappings
-                    aggr_meta.tasks[aggr_task_index] = task
-                    aggr_meta.task_to_task_index[task] = aggr_task_index
-                    aggr_task_index += 1
-
-                # add task_index anyway
-                task_index_to_aggr_task_index[task_index] = aggr_meta.task_to_task_index[task]
-
-            datasets_task_index_to_aggr_task_index[dataset_index] = task_index_to_aggr_task_index
-
-        logging.info("Prepare copy data and videos")
-        datasets_ep_idx_to_aggr_ep_idx = {}
-        datasets_aggr_episode_index_shift = {}
-        aggr_episode_index_shift = 0
-        for dataset_index, meta in enumerate(tqdm.tqdm(all_metadata, desc="Prepare copy data and videos")):
-            ep_idx_to_aggr_ep_idx = {}
-
-            for episode_index in range(meta.total_episodes):
-                aggr_episode_index = episode_index + aggr_episode_index_shift
-                ep_idx_to_aggr_ep_idx[episode_index] = aggr_episode_index
-
-            datasets_ep_idx_to_aggr_ep_idx[dataset_index] = ep_idx_to_aggr_ep_idx
-            datasets_aggr_episode_index_shift[dataset_index] = aggr_episode_index_shift
-
-            # populate episodes
-            for episode_index, episode_dict in meta.episodes.items():
-                aggr_episode_index = episode_index + aggr_episode_index_shift
-                episode_dict["episode_index"] = aggr_episode_index
-                aggr_meta.episodes[aggr_episode_index] = episode_dict
-
-            # populate episodes_stats
-            for episode_index, episode_stats in meta.episodes_stats.items():
-                aggr_episode_index = episode_index + aggr_episode_index_shift
-                aggr_meta.episodes_stats[aggr_episode_index] = episode_stats
-
-            # populate info
-            aggr_meta.info["total_episodes"] += meta.total_episodes
-            aggr_meta.info["total_frames"] += meta.total_frames
-            aggr_meta.info["total_videos"] += len(aggr_meta.video_keys) * meta.total_episodes
-
-            aggr_episode_index_shift += meta.total_episodes
-
-        logging.info("Write meta data")
-        aggr_meta.info["total_tasks"] = len(aggr_meta.tasks)
-        aggr_meta.info["total_chunks"] = aggr_meta.get_episode_chunk(aggr_episode_index_shift - 1)
-        aggr_meta.info["splits"] = {"train": f"0:{aggr_meta.info['total_episodes']}"}
-
-        # create a new episodes jsonl with updated episode_index using write_episode
-        for episode_dict in tqdm.tqdm(aggr_meta.episodes.values(), desc="Write episodes"):
-            write_episode(episode_dict, aggr_meta.root)
-
-        # create a new episode_stats jsonl with updated episode_index using write_episode_stats
-        for episode_index, episode_stats in tqdm.tqdm(
-            aggr_meta.episodes_stats.items(), desc="Write episodes stats"
-        ):
-            write_episode_stats(episode_index, episode_stats, aggr_meta.root)
-
-        # create a new task jsonl with updated episode_index using write_task
-        for task_index, task in tqdm.tqdm(aggr_meta.tasks.items(), desc="Write tasks"):
-            write_task(task_index, task, aggr_meta.root)
-
-        write_info(aggr_meta.info, aggr_meta.root)
-
-        self.datasets_task_index_to_aggr_task_index = datasets_task_index_to_aggr_task_index
-        self.datasets_ep_idx_to_aggr_ep_idx = datasets_ep_idx_to_aggr_ep_idx
-        self.datasets_aggr_episode_index_shift = datasets_aggr_episode_index_shift
-
-        logging.info("Meta data done writing!")
-
-    def run(self, data=None, rank: int = 0, world_size: int = 1):
-        import logging
-        import shutil
-
-        import pandas as pd
-
-        from lerobot.common.datasets.aggregate import get_update_episode_and_task_func
-        from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
-        from lerobot.common.utils.utils import init_logging
-
-        init_logging()
-
-        aggr_meta = LeRobotDatasetMetadata(self.aggr_repo_id)
-        all_metadata = [LeRobotDatasetMetadata(repo_id) for repo_id in self.repo_ids]
-
-        if world_size != len(all_metadata):
-            raise ValueError()
-
-        dataset_index = rank
-        meta = all_metadata[dataset_index]
-        aggr_episode_index_shift = self.datasets_aggr_episode_index_shift[dataset_index]
-
-        logging.info("Copy data")
-        for episode_index in range(meta.total_episodes):
-            aggr_episode_index = self.datasets_ep_idx_to_aggr_ep_idx[dataset_index][episode_index]
-            data_path = meta.root / meta.get_data_file_path(episode_index)
-            aggr_data_path = aggr_meta.root / aggr_meta.get_data_file_path(aggr_episode_index)
-
-            # update episode_index and task_index
-            df = pd.read_parquet(data_path)
-            update_row_func = get_update_episode_and_task_func(
-                aggr_episode_index_shift, self.datasets_task_index_to_aggr_task_index[dataset_index]
-            )
-            df = df.apply(update_row_func, axis=1)
-
-            aggr_data_path.parent.mkdir(parents=True, exist_ok=True)
-            df.to_parquet(aggr_data_path)
-
-        logging.info("Copy videos")
-        for episode_index in range(meta.total_episodes):
-            aggr_episode_index = episode_index + aggr_episode_index_shift
-            for vid_key in meta.video_keys:
-                video_path = meta.root / meta.get_video_file_path(episode_index, vid_key)
-                aggr_video_path = aggr_meta.root / aggr_meta.get_video_file_path(aggr_episode_index, vid_key)
-                aggr_video_path.parent.mkdir(parents=True, exist_ok=True)
-                shutil.copy(video_path, aggr_video_path)
-
-                # copy_command = f"cp {video_path} {aggr_video_path} &"
-                # subprocess.Popen(copy_command, shell=True)
-
-        logging.info("Done!")
-
-
-def make_aggregate_executor(
-    repo_ids, repo_id, job_name, logs_dir, workers, partition, cpus_per_task, mem_per_cpu, slurm=True
-):
-    kwargs = {
-        "pipeline": [
-            AggregateDatasets(repo_ids, repo_id),
-        ],
-        "logging_dir": str(logs_dir / job_name),
-    }
-
-    if slurm:
-        kwargs.update(
-            {
-                "job_name": job_name,
-                "tasks": DROID_SHARDS,
-                "workers": workers,
-                "time": "08:00:00",
-                "partition": partition,
-                "cpus_per_task": cpus_per_task,
-                "sbatch_args": {"mem-per-cpu": mem_per_cpu},
-            }
-        )
-        executor = SlurmPipelineExecutor(**kwargs)
-    else:
-        kwargs.update(
-            {
-                "tasks": DROID_SHARDS,
-                "workers": 1,
-            }
-        )
-        executor = LocalPipelineExecutor(**kwargs)
-
-    return executor
-
-
-def main():
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        help="Repositery identifier on Hugging Face: a community or a user name `/` the name of the dataset, required when push-to-hub is True.",
-    )
-    parser.add_argument(
-        "--logs-dir",
-        type=Path,
-        help="Path to logs directory for `datatrove`.",
-    )
-    parser.add_argument(
-        "--job-name",
-        type=str,
-        default="aggr_droid",
-        help="Job name used in slurm, and name of the directory created inside the provided logs directory.",
-    )
-    parser.add_argument(
-        "--slurm",
-        type=int,
-        default=1,
-        help="Launch over slurm. Use `--slurm 0` to launch sequentially (useful to debug).",
-    )
-    parser.add_argument(
-        "--workers",
-        type=int,
-        default=2048,
-        help="Number of slurm workers. It should be less than the maximum number of shards.",
-    )
-    parser.add_argument(
-        "--partition",
-        type=str,
-        help="Slurm partition. Ideally a CPU partition. No need for GPU partition.",
-    )
-    parser.add_argument(
-        "--cpus-per-task",
-        type=int,
-        default=8,
-        help="Number of cpus that each slurm worker will use.",
-    )
-    parser.add_argument(
-        "--mem-per-cpu",
-        type=str,
-        default="1950M",
-        help="Memory per cpu that each worker will use.",
-    )
-
-    args = parser.parse_args()
-    kwargs = vars(args)
-    kwargs["slurm"] = kwargs.pop("slurm") == 1
-
-    repo_ids = [f"{args.repo_id}_world_{DROID_SHARDS}_rank_{rank}" for rank in range(DROID_SHARDS)]
-    aggregate_executor = make_aggregate_executor(repo_ids, **kwargs)
-    aggregate_executor.run()
-
-
-if __name__ == "__main__":
-    main()

examples/port_datasets/droid_rlds/slurm_port_shards.py

@@ -1,161 +0,0 @@
-import argparse
-from pathlib import Path
-
-from datatrove.executor import LocalPipelineExecutor
-from datatrove.executor.slurm import SlurmPipelineExecutor
-from datatrove.pipeline.base import PipelineStep
-
-from examples.port_datasets.droid_rlds.port_droid import DROID_SHARDS
-from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
-
-
-def validate_shard(repo_id):
-    """Sanity check that ensure meta data can be loaded and all files are present."""
-    meta = LeRobotDatasetMetadata(repo_id)
-
-    if meta.total_episodes == 0:
-        raise ValueError("Number of episodes is 0.")
-
-    for ep_idx in range(meta.total_episodes):
-        data_path = meta.root / meta.get_data_file_path(ep_idx)
-
-        if not data_path.exists():
-            raise ValueError(f"Parquet file is missing in: {data_path}")
-
-        for vid_key in meta.video_keys:
-            vid_path = meta.root / meta.get_video_file_path(ep_idx, vid_key)
-            if not vid_path.exists():
-                raise ValueError(f"Video file is missing in: {vid_path}")
-
-
-class PortDroidShards(PipelineStep):
-    def __init__(
-        self,
-        raw_dir: Path | str,
-        repo_id: str = None,
-    ):
-        super().__init__()
-        self.raw_dir = Path(raw_dir)
-        self.repo_id = repo_id
-
-    def run(self, data=None, rank: int = 0, world_size: int = 1):
-        from datasets.utils.tqdm import disable_progress_bars
-
-        from examples.port_datasets.droid_rlds.port_droid import port_droid
-        from lerobot.common.utils.utils import init_logging
-
-        init_logging()
-        disable_progress_bars()
-
-        shard_repo_id = f"{self.repo_id}_world_{world_size}_rank_{rank}"
-
-        port_droid(
-            self.raw_dir,
-            shard_repo_id,
-            push_to_hub=False,
-            num_shards=world_size,
-            shard_index=rank,
-        )
-
-        validate_shard(shard_repo_id)
-
-
-def make_port_executor(
-    raw_dir, repo_id, job_name, logs_dir, workers, partition, cpus_per_task, mem_per_cpu, slurm=True
-):
-    kwargs = {
-        "pipeline": [
-            PortDroidShards(raw_dir, repo_id),
-        ],
-        "logging_dir": str(logs_dir / job_name),
-    }
-
-    if slurm:
-        kwargs.update(
-            {
-                "job_name": job_name,
-                "tasks": DROID_SHARDS,
-                "workers": workers,
-                "time": "08:00:00",
-                "partition": partition,
-                "cpus_per_task": cpus_per_task,
-                "sbatch_args": {"mem-per-cpu": mem_per_cpu},
-            }
-        )
-        executor = SlurmPipelineExecutor(**kwargs)
-    else:
-        kwargs.update(
-            {
-                "tasks": 1,
-                "workers": 1,
-            }
-        )
-        executor = LocalPipelineExecutor(**kwargs)
-
-    return executor
-
-
-def main():
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--raw-dir",
-        type=Path,
-        required=True,
-        help="Directory containing input raw datasets (e.g. `path/to/dataset` or `path/to/dataset/version).",
-    )
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        help="Repositery identifier on Hugging Face: a community or a user name `/` the name of the dataset, required when push-to-hub is True.",
-    )
-    parser.add_argument(
-        "--logs-dir",
-        type=Path,
-        help="Path to logs directory for `datatrove`.",
-    )
-    parser.add_argument(
-        "--job-name",
-        type=str,
-        default="port_droid",
-        help="Job name used in slurm, and name of the directory created inside the provided logs directory.",
-    )
-    parser.add_argument(
-        "--slurm",
-        type=int,
-        default=1,
-        help="Launch over slurm. Use `--slurm 0` to launch sequentially (useful to debug).",
-    )
-    parser.add_argument(
-        "--workers",
-        type=int,
-        default=2048,
-        help="Number of slurm workers. It should be less than the maximum number of shards.",
-    )
-    parser.add_argument(
-        "--partition",
-        type=str,
-        help="Slurm partition. Ideally a CPU partition. No need for GPU partition.",
-    )
-    parser.add_argument(
-        "--cpus-per-task",
-        type=int,
-        default=8,
-        help="Number of cpus that each slurm worker will use.",
-    )
-    parser.add_argument(
-        "--mem-per-cpu",
-        type=str,
-        default="1950M",
-        help="Memory per cpu that each worker will use.",
-    )
-
-    args = parser.parse_args()
-    kwargs = vars(args)
-    kwargs["slurm"] = kwargs.pop("slurm") == 1
-    port_executor = make_port_executor(**kwargs)
-    port_executor.run()
-
-
-if __name__ == "__main__":
-    main()

examples/port_datasets/droid_rlds/slurm_upload.py

@@ -1,263 +0,0 @@
-import argparse
-import logging
-import os
-from pathlib import Path
-
-from datatrove.executor import LocalPipelineExecutor
-from datatrove.executor.slurm import SlurmPipelineExecutor
-from datatrove.pipeline.base import PipelineStep
-from huggingface_hub import HfApi
-from huggingface_hub.constants import REPOCARD_NAME
-
-from examples.port_datasets.droid_rlds.port_droid import DROID_SHARDS
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDatasetMetadata
-from lerobot.common.datasets.utils import create_lerobot_dataset_card
-from lerobot.common.utils.utils import init_logging
-
-
-class UploadDataset(PipelineStep):
-    def __init__(
-        self,
-        repo_id: str,
-        branch: str | None = None,
-        revision: str | None = None,
-        tags: list | None = None,
-        license: str | None = "apache-2.0",
-        private: bool = False,
-        distant_repo_id: str | None = None,
-        **card_kwargs,
-    ):
-        super().__init__()
-        self.repo_id = repo_id
-        self.distant_repo_id = self.repo_id if distant_repo_id is None else distant_repo_id
-        self.branch = branch
-        self.tags = tags
-        self.license = license
-        self.private = private
-        self.card_kwargs = card_kwargs
-        self.revision = revision if revision else CODEBASE_VERSION
-
-        if os.environ.get("HF_HUB_ENABLE_HF_TRANSFER", "0") != "1":
-            logging.warning(
-                'HF_HUB_ENABLE_HF_TRANSFER is not set to "1". Install hf_transfer and set the env '
-                "variable for faster uploads:\npip install hf-transfer\nexport HF_HUB_ENABLE_HF_TRANSFER=1"
-            )
-
-        self.create_repo()
-
-    def create_repo(self):
-        logging.info(f"Loading meta data from {self.repo_id}...")
-        meta = LeRobotDatasetMetadata(self.repo_id)
-
-        logging.info(f"Creating repo {self.distant_repo_id}...")
-        hub_api = HfApi()
-        hub_api.create_repo(
-            repo_id=self.distant_repo_id,
-            private=self.private,
-            repo_type="dataset",
-            exist_ok=True,
-        )
-        if self.branch:
-            hub_api.create_branch(
-                repo_id=self.distant_repo_id,
-                branch=self.branch,
-                revision=self.revision,
-                repo_type="dataset",
-                exist_ok=True,
-            )
-
-        if not hub_api.file_exists(
-            self.distant_repo_id, REPOCARD_NAME, repo_type="dataset", revision=self.branch
-        ):
-            card = create_lerobot_dataset_card(
-                tags=self.tags, dataset_info=meta.info, license=self.license, **self.card_kwargs
-            )
-            card.push_to_hub(repo_id=self.distant_repo_id, repo_type="dataset", revision=self.branch)
-
-        def list_files_recursively(directory):
-            base_path = Path(directory)
-            return [str(file.relative_to(base_path)) for file in base_path.rglob("*") if file.is_file()]
-
-        logging.info(f"Listing all local files from {self.repo_id}...")
-        self.file_paths = list_files_recursively(meta.root)
-        self.file_paths = sorted(self.file_paths)
-
-    def create_chunks(self, lst, n):
-        from itertools import islice
-
-        it = iter(lst)
-        return [list(islice(it, size)) for size in [len(lst) // n + (i < len(lst) % n) for i in range(n)]]
-
-    def create_commits(self, additions):
-        import logging
-        import math
-        import random
-        import time
-
-        from huggingface_hub import create_commit
-        from huggingface_hub.utils import HfHubHTTPError
-
-        FILES_BETWEEN_COMMITS = 10  # noqa: N806
-        BASE_DELAY = 0.1  # noqa: N806
-        MAX_RETRIES = 12  # noqa: N806
-
-        # Split the files into smaller chunks for faster commit
-        # and avoiding "A commit has happened since" error
-        num_chunks = math.ceil(len(additions) / FILES_BETWEEN_COMMITS)
-        chunks = self.create_chunks(additions, num_chunks)
-
-        for chunk in chunks:
-            retries = 0
-            while True:
-                try:
-                    create_commit(
-                        self.distant_repo_id,
-                        repo_type="dataset",
-                        operations=chunk,
-                        commit_message=f"DataTrove upload ({len(chunk)} files)",
-                        revision=self.branch,
-                    )
-                    # TODO: every 100 chunks super_squach_commits()
-                    logging.info("create_commit completed!")
-                    break
-                except HfHubHTTPError as e:
-                    if "A commit has happened since" in e.server_message:
-                        if retries >= MAX_RETRIES:
-                            logging.error(f"Failed to create commit after {MAX_RETRIES=}. Giving up.")
-                            raise e
-                        logging.info("Commit creation race condition issue. Waiting...")
-                        time.sleep(BASE_DELAY * 2**retries + random.uniform(0, 2))
-                        retries += 1
-                    else:
-                        raise e
-
-    def run(self, data=None, rank: int = 0, world_size: int = 1):
-        import logging
-
-        from datasets.utils.tqdm import disable_progress_bars
-        from huggingface_hub import CommitOperationAdd, preupload_lfs_files
-
-        from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
-        from lerobot.common.utils.utils import init_logging
-
-        init_logging()
-        disable_progress_bars()
-
-        chunks = self.create_chunks(self.file_paths, world_size)
-        file_paths = chunks[rank]
-
-        if len(file_paths) == 0:
-            raise ValueError(file_paths)
-
-        logging.info("Pre-uploading LFS files...")
-        for i, path in enumerate(file_paths):
-            logging.info(f"{i}: {path}")
-
-        meta = LeRobotDatasetMetadata(self.repo_id)
-        additions = [
-            CommitOperationAdd(path_in_repo=path, path_or_fileobj=meta.root / path) for path in file_paths
-        ]
-        preupload_lfs_files(
-            repo_id=self.distant_repo_id, repo_type="dataset", additions=additions, revision=self.branch
-        )
-
-        logging.info("Creating commits...")
-        self.create_commits(additions)
-        logging.info("Done!")
-
-
-def make_upload_executor(
-    repo_id, job_name, logs_dir, workers, partition, cpus_per_task, mem_per_cpu, slurm=True
-):
-    kwargs = {
-        "pipeline": [
-            UploadDataset(repo_id),
-        ],
-        "logging_dir": str(logs_dir / job_name),
-    }
-
-    if slurm:
-        kwargs.update(
-            {
-                "job_name": job_name,
-                "tasks": DROID_SHARDS,
-                "workers": workers,
-                "time": "08:00:00",
-                "partition": partition,
-                "cpus_per_task": cpus_per_task,
-                "sbatch_args": {"mem-per-cpu": mem_per_cpu},
-            }
-        )
-        executor = SlurmPipelineExecutor(**kwargs)
-    else:
-        kwargs.update(
-            {
-                "tasks": DROID_SHARDS,
-                "workers": 1,
-            }
-        )
-        executor = LocalPipelineExecutor(**kwargs)
-
-    return executor
-
-
-def main():
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        help="Repositery identifier on Hugging Face: a community or a user name `/` the name of the dataset, required when push-to-hub is True.",
-    )
-    parser.add_argument(
-        "--logs-dir",
-        type=Path,
-        help="Path to logs directory for `datatrove`.",
-    )
-    parser.add_argument(
-        "--job-name",
-        type=str,
-        default="upload_droid",
-        help="Job name used in slurm, and name of the directory created inside the provided logs directory.",
-    )
-    parser.add_argument(
-        "--slurm",
-        type=int,
-        default=1,
-        help="Launch over slurm. Use `--slurm 0` to launch sequentially (useful to debug).",
-    )
-    parser.add_argument(
-        "--workers",
-        type=int,
-        default=50,
-        help="Number of slurm workers. It should be less than the maximum number of shards.",
-    )
-    parser.add_argument(
-        "--partition",
-        type=str,
-        help="Slurm partition. Ideally a CPU partition. No need for GPU partition.",
-    )
-    parser.add_argument(
-        "--cpus-per-task",
-        type=int,
-        default=8,
-        help="Number of cpus that each slurm worker will use.",
-    )
-    parser.add_argument(
-        "--mem-per-cpu",
-        type=str,
-        default="1950M",
-        help="Memory per cpu that each worker will use.",
-    )
-
-    init_logging()
-
-    args = parser.parse_args()
-    kwargs = vars(args)
-    kwargs["slurm"] = kwargs.pop("slurm") == 1
-    upload_executor = make_upload_executor(**kwargs)
-    upload_executor.run()
-
-
-if __name__ == "__main__":
-    main()

lerobot/common/datasets/aggregate.py

@@ -1,175 +0,0 @@
-import logging
-import shutil
-
-import pandas as pd
-import tqdm
-
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
-from lerobot.common.datasets.utils import write_episode, write_episode_stats, write_info, write_task
-from lerobot.common.utils.utils import init_logging
-
-
-def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
-    # validate same fps, robot_type, features
-
-    fps = all_metadata[0].fps
-    robot_type = all_metadata[0].robot_type
-    features = all_metadata[0].features
-
-    for meta in tqdm.tqdm(all_metadata, desc="Validate all meta data"):
-        if fps != meta.fps:
-            raise ValueError(f"Same fps is expected, but got fps={meta.fps} instead of {fps}.")
-        if robot_type != meta.robot_type:
-            raise ValueError(
-                f"Same robot_type is expected, but got robot_type={meta.robot_type} instead of {robot_type}."
-            )
-        if features != meta.features:
-            raise ValueError(
-                f"Same features is expected, but got features={meta.features} instead of {features}."
-            )
-
-    return fps, robot_type, features
-
-
-def get_update_episode_and_task_func(episode_index_to_add, task_index_to_global_task_index):
-    def _update(row):
-        row["episode_index"] = row["episode_index"] + episode_index_to_add
-        row["task_index"] = task_index_to_global_task_index[row["task_index"]]
-        return row
-
-    return _update
-
-
-def aggregate_datasets(repo_ids: list[str], aggr_repo_id: str, aggr_root=None):
-    logging.info("Start aggregate_datasets")
-
-    all_metadata = [LeRobotDatasetMetadata(repo_id) for repo_id in repo_ids]
-
-    fps, robot_type, features = validate_all_metadata(all_metadata)
-
-    # Create resulting dataset folder
-    aggr_meta = LeRobotDatasetMetadata.create(
-        repo_id=aggr_repo_id,
-        fps=fps,
-        robot_type=robot_type,
-        features=features,
-        root=aggr_root,
-    )
-
-    logging.info("Find all tasks")
-    # find all tasks, deduplicate them, create new task indices for each dataset
-    # indexed by dataset index
-    datasets_task_index_to_aggr_task_index = {}
-    aggr_task_index = 0
-    for dataset_index, meta in enumerate(tqdm.tqdm(all_metadata, desc="Find all tasks")):
-        task_index_to_aggr_task_index = {}
-
-        for task_index, task in meta.tasks.items():
-            if task not in aggr_meta.task_to_task_index:
-                # add the task to aggr tasks mappings
-                aggr_meta.tasks[aggr_task_index] = task
-                aggr_meta.task_to_task_index[task] = aggr_task_index
-                aggr_task_index += 1
-
-            # add task_index anyway
-            task_index_to_aggr_task_index[task_index] = aggr_meta.task_to_task_index[task]
-
-        datasets_task_index_to_aggr_task_index[dataset_index] = task_index_to_aggr_task_index
-
-    logging.info("Copy data and videos")
-    aggr_episode_index_shift = 0
-    for dataset_index, meta in enumerate(tqdm.tqdm(all_metadata, desc="Copy data and videos")):
-        # cp data
-        for episode_index in range(meta.total_episodes):
-            aggr_episode_index = episode_index + aggr_episode_index_shift
-            data_path = meta.root / meta.get_data_file_path(episode_index)
-            aggr_data_path = aggr_meta.root / aggr_meta.get_data_file_path(aggr_episode_index)
-
-            # update episode_index and task_index
-            df = pd.read_parquet(data_path)
-            update_row_func = get_update_episode_and_task_func(
-                aggr_episode_index_shift, datasets_task_index_to_aggr_task_index[dataset_index]
-            )
-            df = df.apply(update_row_func, axis=1)
-
-            aggr_data_path.parent.mkdir(parents=True, exist_ok=True)
-            df.to_parquet(aggr_data_path)
-
-        # cp videos
-        for episode_index in range(meta.total_episodes):
-            aggr_episode_index = episode_index + aggr_episode_index_shift
-            for vid_key in meta.video_keys:
-                video_path = meta.root / meta.get_video_file_path(episode_index, vid_key)
-                aggr_video_path = aggr_meta.root / aggr_meta.get_video_file_path(aggr_episode_index, vid_key)
-                aggr_video_path.parent.mkdir(parents=True, exist_ok=True)
-                shutil.copy(video_path, aggr_video_path)
-
-                # copy_command = f"cp {video_path} {aggr_video_path} &"
-                # subprocess.Popen(copy_command, shell=True)
-
-        # populate episodes
-        for episode_index, episode_dict in meta.episodes.items():
-            aggr_episode_index = episode_index + aggr_episode_index_shift
-            episode_dict["episode_index"] = aggr_episode_index
-            aggr_meta.episodes[aggr_episode_index] = episode_dict
-
-        # populate episodes_stats
-        for episode_index, episode_stats in meta.episodes_stats.items():
-            aggr_episode_index = episode_index + aggr_episode_index_shift
-            aggr_meta.episodes_stats[aggr_episode_index] = episode_stats
-
-        # populate info
-        aggr_meta.info["total_episodes"] += meta.total_episodes
-        aggr_meta.info["total_frames"] += meta.total_frames
-        aggr_meta.info["total_videos"] += len(aggr_meta.video_keys) * meta.total_episodes
-
-        aggr_episode_index_shift += meta.total_episodes
-
-    logging.info("write meta data")
-
-    aggr_meta.info["total_chunks"] = aggr_meta.get_episode_chunk(aggr_episode_index_shift - 1)
-    aggr_meta.info["splits"] = {"train": f"0:{aggr_meta.info['total_episodes']}"}
-
-    # create a new episodes jsonl with updated episode_index using write_episode
-    for episode_dict in aggr_meta.episodes.values():
-        write_episode(episode_dict, aggr_meta.root)
-
-    # create a new episode_stats jsonl with updated episode_index using write_episode_stats
-    for episode_index, episode_stats in aggr_meta.episodes_stats.items():
-        write_episode_stats(episode_index, episode_stats, aggr_meta.root)
-
-    # create a new task jsonl with updated episode_index using write_task
-    for task_index, task in aggr_meta.tasks.items():
-        write_task(task_index, task, aggr_meta.root)
-
-    write_info(aggr_meta.info, aggr_meta.root)
-
-
-if __name__ == "__main__":
-    init_logging()
-    repo_id = "cadene/droid"
-    aggr_repo_id = "cadene/droid"
-    datetime = "2025-02-22_11-23-54"
-
-    # root = Path(f"/tmp/{repo_id}")
-    # if root.exists():
-    #     shutil.rmtree(root)
-    root = None
-
-    # all_metadata = [LeRobotDatasetMetadata(f"{repo_id}_{datetime}_world_2048_rank_{rank}") for rank in range(2048)]
-
-    # aggregate_datasets(
-    #     all_metadata,
-    #     aggr_repo_id,
-    #     root=root,
-    # )
-
-    aggr_dataset = LeRobotDataset(
-        repo_id=aggr_repo_id,
-        root=root,
-    )
-    aggr_dataset.push_to_hub(tags=["openx"])
-
-    # for meta in all_metadata:
-    #     dataset = LeRobotDataset(repo_id=meta.repo_id, root=meta.root)
-    #     dataset.push_to_hub(tags=["openx"])

lerobot/common/datasets/lerobot_dataset.py

@@ -73,7 +73,7 @@ from lerobot.common.datasets.video_utils import (
 )
 from lerobot.common.robot_devices.robots.utils import Robot
 
-CODEBASE_VERSION = "v3.0"
+CODEBASE_VERSION = "v2.1"
 
 
 class LeRobotDatasetMetadata:
@@ -616,8 +616,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
         """hf_dataset contains all the observations, states, actions, rewards, etc."""
         if self.episodes is None:
             path = str(self.root / "data")
-            # TODO(rcadene): load_dataset convert parquet to arrow.
-            # set num_proc to accelerate this conversion
             hf_dataset = load_dataset("parquet", data_dir=path, split="train")
         else:
             files = [str(self.root / self.meta.get_data_file_path(ep_idx)) for ep_idx in self.episodes]

lerobot/common/datasets/v30/convert_dataset_v21_to_v30.py

@@ -1,137 +0,0 @@
-"""
-This script will help you convert any LeRobot dataset already pushed to the hub from codebase version 2.1 to
-3.0. It will:
-
-- Generate per-episodes stats and writes them in `episodes_stats.jsonl`
-- Check consistency between these new stats and the old ones.
-- Remove the deprecated `stats.json`.
-- Update codebase_version in `info.json`.
-- Push this new version to the hub on the 'main' branch and tags it with "v2.1".
-
-Usage:
-
-```bash
-python lerobot/common/datasets/v30/convert_dataset_v21_to_v30.py \
-    --repo-id=lerobot/pusht
-```
-
-"""
-
-import argparse
-import logging
-
-from datasets import Dataset
-from huggingface_hub import snapshot_download
-
-from lerobot.common.constants import HF_LEROBOT_HOME
-from lerobot.common.datasets.utils import (
-    load_episodes_stats,
-)
-
-V21 = "v2.1"
-
-
-class SuppressWarnings:
-    def __enter__(self):
-        self.previous_level = logging.getLogger().getEffectiveLevel()
-        logging.getLogger().setLevel(logging.ERROR)
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        logging.getLogger().setLevel(self.previous_level)
-
-
-def convert_dataset(
-    repo_id: str,
-    branch: str | None = None,
-    num_workers: int = 4,
-):
-    root = HF_LEROBOT_HOME / repo_id
-    snapshot_download(
-        repo_id,
-        repo_type="dataset",
-        revision=V21,
-        local_dir=root,
-    )
-
-    # Concatenate videos
-
-    # Create
-
-    """
-    -------------------------
-    OLD
-    data/chunk-000/episode_000000.parquet
-
-    NEW
-    data/chunk-000/file_000.parquet
-    -------------------------
-    OLD
-    videos/chunk-000/CAMERA/episode_000000.mp4
-
-    NEW
-    videos/chunk-000/file_000.mp4
-    -------------------------
-    OLD
-    episodes.jsonl
-    {"episode_index": 1, "tasks": ["Put the blue block in the green bowl"], "length": 266}
-
-    NEW
-    meta/episodes/chunk-000/episodes_000.parquet
-    episode_index | video_chunk_index | video_file_index | data_chunk_index | data_file_index | tasks | length
-    -------------------------
-    OLD
-    tasks.jsonl
-    {"task_index": 1, "task": "Put the blue block in the green bowl"}
-
-    NEW
-    meta/tasks/chunk-000/file_000.parquet
-    task_index | task
-    -------------------------
-    OLD
-    episodes_stats.jsonl
-
-    NEW
-    meta/episodes_stats/chunk-000/file_000.parquet
-    episode_index | mean | std | min | max
-    -------------------------
-    UPDATE
-    meta/info.json
-    -------------------------
-    """
-
-    new_root = HF_LEROBOT_HOME / f"{repo_id}_v30"
-    new_root.mkdir(parents=True, exist_ok=True)
-
-    episodes_stats = load_episodes_stats(root)
-    hf_dataset = Dataset.from_dict(episodes_stats)  # noqa: F841
-
-    meta_ep_st_ch = new_root / "meta/episodes_stats/chunk-000"
-    meta_ep_st_ch.mkdir(parents=True, exist_ok=True)
-
-    # hf_dataset.to_parquet(meta_ep_st_ch / 'file_000.parquet')
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        required=True,
-        help="Repository identifier on Hugging Face: a community or a user name `/` the name of the dataset "
-        "(e.g. `lerobot/pusht`, `cadene/aloha_sim_insertion_human`).",
-    )
-    parser.add_argument(
-        "--branch",
-        type=str,
-        default=None,
-        help="Repo branch to push your dataset. Defaults to the main branch.",
-    )
-    parser.add_argument(
-        "--num-workers",
-        type=int,
-        default=4,
-        help="Number of workers for parallelizing stats compute. Defaults to 4.",
-    )
-
-    args = parser.parse_args()
-    convert_dataset(**vars(args))

lerobot/common/datasets/video_utils.py

@@ -136,7 +136,7 @@ def encode_video_frames(
     g: int | None = 2,
     crf: int | None = 30,
     fast_decode: int = 0,
-    log_level: str | None = "quiet",
+    log_level: str | None = "error",
     overwrite: bool = False,
 ) -> None:
     """More info on ffmpeg arguments tuning on `benchmark/video/README.md`"""

lerobot/common/policies/auto/configuration_auto.py

@@ -0,0 +1,409 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib
+import logging
+import os
+from collections import OrderedDict
+from pathlib import Path
+from typing import Any, Dict, Optional, Type, Union
+
+from lerobot.common.policies.pretrained import PreTrainedPolicy
+from lerobot.configs.policies import PreTrainedConfig
+
+logger = logging.getLogger(__name__)
+
+# Constants
+IMPORT_PATHS = ["lerobot.common.policies.{0}.configuration_{0}"]
+
+POLICY_IMPORT_PATHS = ["lerobot.common.policies.{0}.modeling_{0}"]
+
+
+def policy_type_to_module_name(policy_type: str) -> str:
+    """
+    Convert policy type to module name format.
+
+    Args:
+        policy_type: The policy type identifier (e.g. 'lerobot/vqbet-pusht')
+
+    Returns:
+        str: Normalized module name (e.g. 'vqbet')
+
+    Examples:
+        >>> policy_type_to_module_name("lerobot/vqbet-pusht")
+        'vqbet'
+    """
+    # TODO(Steven): This is a temporary solution, we should have a more robust way to handle this
+    return policy_type.replace("lerobot/", "").replace("-", "_").replace("_", "").replace("pusht", "")
+
+
+class _LazyPolicyConfigMapping(OrderedDict):
+    def __init__(self, mapping: Dict[str, str]):
+        self._mapping = mapping
+        self._extra_content: Dict[str, Any] = {}
+        self._modules: Dict[str, Any] = {}
+
+    def __getitem__(self, key: str) -> Any:
+        if key in self._extra_content:
+            return self._extra_content[key]
+        if key not in self._mapping:
+            raise KeyError(f"Policy type '{key}' not found in mapping")
+
+        value = self._mapping[key]
+        module_name = policy_type_to_module_name(key)
+
+        for import_path in IMPORT_PATHS:
+            try:
+                if key not in self._modules:
+                    self._modules[key] = importlib.import_module(import_path.format(module_name))
+                    logger.debug(f"Config module: {module_name} imported")
+                if hasattr(self._modules[key], value):
+                    return getattr(self._modules[key], value)
+            except ImportError:
+                continue
+
+        raise ImportError(f"Could not find configuration class {value} for policy type {key}")
+
+    def keys(self):
+        return list(self._mapping.keys()) + list(self._extra_content.keys())
+
+    def values(self):
+        return [self[k] for k in self._mapping] + list(self._extra_content.values())
+
+    def items(self):
+        return [(k, self[k]) for k in self._mapping] + list(self._extra_content.items())
+
+    def __iter__(self):
+        return iter(list(self._mapping.keys()) + list(self._extra_content.keys()))
+
+    def __contains__(self, item):
+        return item in self._mapping or item in self._extra_content
+
+    def register(self, key, value, exist_ok=False):
+        """
+        Register a new configuration in this mapping.
+        """
+        if key in self._mapping and not exist_ok:
+            raise ValueError(f"'{key}' is already used by a Policy Config, pick another name.")
+        self._extra_content[key] = value
+
+
+POLICY_CONFIG_NAMES_MAPPING = OrderedDict(
+    [
+        ("vqbet", "VQBeTConfig"),
+        ("lerobot/vqbet_pusht", "VQBeTConfig"),
+    ]
+)
+
+POLICY_CONFIG_MAPPING = _LazyPolicyConfigMapping(POLICY_CONFIG_NAMES_MAPPING)
+
+
+class _LazyPolicyMapping(OrderedDict):
+    """
+    A dictionary that lazily loads its values when they are requested.
+    """
+
+    def __init__(self, mapping: Dict[str, str]):
+        self._mapping = mapping
+        self._extra_content: Dict[str, Type[PreTrainedPolicy]] = {}
+        self._modules: Dict[str, Any] = {}
+        self._config_mapping: Dict[Type[PreTrainedConfig], Type[PreTrainedPolicy]] = {}
+        self._initialized_types: set[str] = set()
+
+    def _lazy_init_for_type(self, policy_type: str) -> None:
+        """Lazily initialize mappings for a policy type if not already done."""
+        if policy_type not in self._initialized_types:
+            try:
+                config_class = POLICY_CONFIG_MAPPING[policy_type]
+                self._config_mapping[config_class] = self[policy_type]
+                self._initialized_types.add(policy_type)
+            except (ImportError, AttributeError, KeyError) as e:
+                logger.warning(f"Could not automatically map config for policy type {policy_type}: {str(e)}")
+
+    def __getitem__(self, key: str) -> Type[PreTrainedPolicy]:
+        """Get a policy class by key with lazy loading."""
+        if key in self._extra_content:
+            return self._extra_content[key]
+        if key not in self._mapping:
+            raise KeyError(f"Policy type '{key}' not found in mapping")
+
+        value = self._mapping[key]
+        module_name = policy_type_to_module_name(key)
+
+        for import_path in POLICY_IMPORT_PATHS:
+            try:
+                if key not in self._modules:
+                    self._modules[key] = importlib.import_module(import_path.format(module_name))
+                    logger.debug(
+                        f"Policy module: {module_name} imported from {import_path.format(module_name)}"
+                    )
+                if hasattr(self._modules[key], value):
+                    return getattr(self._modules[key], value)
+            except ImportError:
+                continue
+
+        raise ImportError(
+            f"Could not find policy class {value} for policy type {key}. "
+            f"Tried paths: {[p.format(module_name) for p in POLICY_IMPORT_PATHS]}"
+        )
+
+    def register(
+        self,
+        key: str,
+        value: Type[PreTrainedPolicy],
+        config_class: Type[PreTrainedConfig],
+        exist_ok: bool = False,
+    ) -> None:
+        """Register a new policy class with its configuration class."""
+        if not isinstance(key, str):
+            raise TypeError(f"Key must be a string, got {type(key)}")
+        if not issubclass(value, PreTrainedPolicy):
+            raise TypeError(f"Value must be a PreTrainedPolicy subclass, got {type(value)}")
+        if not issubclass(config_class, PreTrainedConfig):
+            raise TypeError(f"Config class must be a PreTrainedConfig subclass, got {type(config_class)}")
+
+        if key in self._mapping and not exist_ok:
+            raise ValueError(f"'{key}' is already used by a Policy, pick another name.")
+        self._extra_content[key] = value
+        self._config_mapping[config_class] = value
+
+    def get_policy_for_config(self, config_class: Type[PreTrainedConfig]) -> Type[PreTrainedPolicy]:
+        """Get the policy class associated with a config class."""
+        # First check direct config class mapping
+        if config_class in self._config_mapping:
+            return self._config_mapping[config_class]
+
+        # Try to find by policy type
+        try:
+            policy_type = config_class.get_type_str()
+            # Check extra content first
+            if policy_type in self._extra_content:
+                return self._extra_content[policy_type]
+
+            # Then check standard mapping
+            if policy_type in self._mapping:
+                self._lazy_init_for_type(policy_type)
+                if config_class in self._config_mapping:
+                    return self._config_mapping[config_class]
+                return self[policy_type]
+        except AttributeError:
+            pass
+
+        raise ValueError(
+            f"No policy class found for config class {config_class.__name__}. "
+            f"Available types: {list(self._mapping.keys()) + list(self._extra_content.keys())}"
+        )
+
+
+POLICY_NAMES_MAPPING = OrderedDict(
+    [
+        ("vqbet", "VQBeTPolicy"),
+        ("lerobot/vqbet_pusht", "VQBeTPolicy"),
+    ]
+)
+
+POLICY_MAPPING = _LazyPolicyMapping(POLICY_NAMES_MAPPING)
+
+
+class AutoPolicyConfig:
+    """
+    Factory class for automatically loading policy configurations.
+
+    This class provides methods to:
+    - Load pre-trained policy configurations from local files or the Hub
+    - Register new policy types dynamically
+    - Create policy configurations for specific policy types
+    """
+
+    def __init__(self):
+        raise OSError("AutoPolicyConfig not meant to be instantiated directly")
+
+    @classmethod
+    def for_policy(cls, policy_type: str, *args, **kwargs) -> PreTrainedConfig:
+        """Create a new configuration instance for the specified policy type."""
+        if policy_type in POLICY_CONFIG_MAPPING:
+            config_class = POLICY_CONFIG_MAPPING[policy_type]
+            return config_class(*args, **kwargs)
+        raise ValueError(
+            f"Unrecognized policy identifier: {policy_type}. Should contain one of {', '.join(POLICY_CONFIG_MAPPING.keys())}"
+        )
+
+    @staticmethod
+    def register(policy_type, config, exist_ok=False):
+        """
+        Register a new configuration for this class.
+
+        Args:
+            policy_type (`str`): The policy type like "act" or "pi0".
+            config ([`PreTrainedConfig`]): The config to register.
+        """
+        if issubclass(config, PreTrainedConfig) and config.get_type_str() != policy_type:
+            raise ValueError(
+                "The config you are passing has a `policy_type` attribute that is not consistent with the policy type "
+                f"you passed (config has {config.type} and you passed {policy_type}. Fix one of those so they "
+                "match!"
+            )
+        POLICY_CONFIG_MAPPING.register(policy_type, config, exist_ok=exist_ok)
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_policy_config_name_or_path: Union[str, Path], **kwargs
+    ) -> PreTrainedConfig:
+        """
+        Instantiate a PreTrainedConfig from a pre-trained policy configuration.
+
+        Args:
+            pretrained_policy_config_name_or_path (`str` or `Path`):
+                Can be either:
+                    - A string with the `policy_type` of a pre-trained policy configuration listed on
+                      the Hub or locally (e.g., 'act')
+                    - A path to a `directory` containing a configuration file saved
+                      using [`~PreTrainedConfig.save_pretrained`].
+                    - A path or url to a saved configuration JSON `file`.
+            **kwargs: Additional kwargs passed to PreTrainedConfig.from_pretrained()
+
+        Returns:
+            [`PreTrainedConfig`]: The configuration object instantiated from that pre-trained policy config.
+        """
+        if os.path.isdir(pretrained_policy_config_name_or_path):
+            # Load from local directory
+            config_dict = PreTrainedConfig.from_pretrained(pretrained_policy_config_name_or_path, **kwargs)
+            policy_type = config_dict.type
+        elif os.path.isfile(pretrained_policy_config_name_or_path):
+            # Load from local file
+            config_dict = PreTrainedConfig.from_pretrained(pretrained_policy_config_name_or_path, **kwargs)
+            policy_type = config_dict.type
+        else:
+            # Assume it's a policy_type identifier
+            policy_type = pretrained_policy_config_name_or_path
+
+        if policy_type not in POLICY_CONFIG_MAPPING:
+            raise ValueError(
+                f"Unrecognized policy type {policy_type}. "
+                f"Should be one of {', '.join(POLICY_CONFIG_MAPPING.keys())}"
+            )
+
+        config_class = POLICY_CONFIG_MAPPING[policy_type]
+        return config_class.from_pretrained(pretrained_policy_config_name_or_path, **kwargs)
+
+
+class AutoPolicy:
+    """
+    Factory class that allows instantiating policy models from configurations.
+
+    This class provides methods to:
+    - Load pre-trained policies from configurations
+    - Register new policy types dynamically
+    - Create policy instances for specific configurations
+    """
+
+    def __init__(self):
+        raise OSError("AutoPolicy not meant to be instantiated directly")
+
+    @classmethod
+    def from_config(cls, config: PreTrainedConfig, **kwargs) -> PreTrainedPolicy:
+        """Instantiate a policy from a configuration."""
+        policy_class = POLICY_MAPPING.get_policy_for_config(type(config))
+        return policy_class(config, **kwargs)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_policy_name_or_path: Union[str, Path],
+        *,
+        config: Optional[PreTrainedConfig] = None,
+        **kwargs,
+    ) -> PreTrainedPolicy:
+        """
+        Instantiate a pre-trained policy from a configuration.
+
+        Args:
+            pretrained_policy_name_or_path: Path to pretrained weights or model identifier
+            config: Optional configuration for the policy
+            **kwargs: Additional arguments to pass to from_pretrained()
+        """
+        if config is None:
+            config = AutoPolicyConfig.from_pretrained(pretrained_policy_name_or_path)
+
+        if isinstance(config, str):
+            config = AutoPolicyConfig.from_pretrained(config)
+
+        policy_class = POLICY_MAPPING.get_policy_for_config(config)
+        return policy_class.from_pretrained(pretrained_policy_name_or_path, config=config, **kwargs)
+
+    @staticmethod
+    def register(
+        config_class: Type[PreTrainedConfig], policy_class: Type[PreTrainedPolicy], exist_ok: bool = False
+    ):
+        """
+        Register a new policy class for a configuration class.
+
+        Args:
+            config_class: The configuration class
+            policy_class: The policy class to register
+            exist_ok: Whether to allow overwriting existing registrations
+        """
+        POLICY_MAPPING.register(config_class.get_type_str(), policy_class, config_class, exist_ok=exist_ok)
+
+
+def main():
+    """Test the AutoPolicy and AutoPolicyConfig functionality."""
+
+    def test_error_cases():
+        """Test error handling"""
+        try:
+            AutoPolicyConfig()
+        except OSError as e:
+            assert "not meant to be instantiated directly" in str(e)
+        try:
+            AutoPolicy()
+        except OSError as e:
+            assert "not meant to be instantiated directly" in str(e)
+
+        # try:
+        #     AutoPolicy.from_config("invalid_config")
+        # except ValueError as e:
+        #     assert "Unrecognized policy identifier" in str(e)
+
+    logging.basicConfig(level=logging.DEBUG)
+
+    # Test built-in policy loading
+    # config = AutoPolicyConfig.from_pretrained("lerobot/vqbet_pusht")
+    config = AutoPolicyConfig.for_policy("vqbet")
+    policy = AutoPolicy.from_config(config)
+
+    from lerobot.common.policies.vqbet.configuration_vqbet import VQBeTConfig
+    from lerobot.common.policies.vqbet.modeling_vqbet import VQBeTPolicy
+
+    assert isinstance(config, VQBeTConfig)
+    assert isinstance(policy, VQBeTPolicy)
+
+    # Test policy registration
+    from lerobot.common.policies.tdmpc.configuration_tdmpc import TDMPCConfig
+    from lerobot.common.policies.tdmpc.modeling_tdmpc import TDMPCPolicy
+
+    AutoPolicyConfig.register("tdmpc", TDMPCConfig)
+    AutoPolicy.register(TDMPCConfig, TDMPCPolicy)
+
+    my_new_config = AutoPolicyConfig.for_policy("tdmpc")
+    my_new_policy = AutoPolicy.from_config(my_new_config)
+    assert isinstance(my_new_config, TDMPCConfig)
+    assert isinstance(my_new_policy, TDMPCPolicy)
+
+    # Run error case tests
+    test_error_cases()
+
+
+if __name__ == "__main__":
+    main()

lerobot/common/policies/pi0/modeling_pi0.py

@@ -313,7 +313,7 @@ class PI0Policy(PreTrainedPolicy):
         state = self.prepare_state(batch)
         lang_tokens, lang_masks = self.prepare_language(batch)
         actions = self.prepare_action(batch)
-        actions_is_pad = batch.get("actions_id_pad")
+        actions_is_pad = batch.get("actions_is_pad")
 
         loss_dict = {}
         losses = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions, noise, time)

lerobot/common/robot_devices/cameras/utils.py

@@ -31,7 +31,7 @@ def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> list[C
 
             cameras[key] = IntelRealSenseCamera(cfg)
         else:
-            raise ValueError(f"The motor type '{cfg.type}' is not valid.")
+            raise ValueError(f"The camera type '{cfg.type}' is not valid.")
 
     return cameras
 

lerobot/common/robot_devices/robots/mobile_manipulator.py

@@ -392,21 +392,19 @@ class MobileManipulator:
         for name in self.leader_arms:
             pos = self.leader_arms[name].read("Present_Position")
             pos_tensor = torch.from_numpy(pos).float()
-            # Instead of pos_tensor.item(), use tolist() to convert the entire tensor to a list
             arm_positions.extend(pos_tensor.tolist())
 
-        # (The rest of your code for generating wheel commands remains unchanged)
-        x_cmd = 0.0  # m/s forward/backward
-        y_cmd = 0.0  # m/s lateral
+        y_cmd = 0.0  # m/s forward/backward
+        x_cmd = 0.0  # m/s lateral
         theta_cmd = 0.0  # deg/s rotation
         if self.pressed_keys["forward"]:
-            x_cmd += xy_speed
-        if self.pressed_keys["backward"]:
-            x_cmd -= xy_speed
-        if self.pressed_keys["left"]:
             y_cmd += xy_speed
-        if self.pressed_keys["right"]:
+        if self.pressed_keys["backward"]:
             y_cmd -= xy_speed
+        if self.pressed_keys["left"]:
+            x_cmd += xy_speed
+        if self.pressed_keys["right"]:
+            x_cmd -= xy_speed
         if self.pressed_keys["rotate_left"]:
             theta_cmd += theta_speed
         if self.pressed_keys["rotate_right"]:
@@ -584,8 +582,8 @@ class MobileManipulator:
         # Create the body velocity vector [x, y, theta_rad].
         velocity_vector = np.array([x_cmd, y_cmd, theta_rad])
 
-        # Define the wheel mounting angles with a -90° offset.
-        angles = np.radians(np.array([240, 120, 0]) - 90)
+        # Define the wheel mounting angles (defined from y axis cw)
+        angles = np.radians(np.array([300, 180, 60]))
         # Build the kinematic matrix: each row maps body velocities to a wheel’s linear speed.
         # The third column (base_radius) accounts for the effect of rotation.
         m = np.array([[np.cos(a), np.sin(a), base_radius] for a in angles])
@@ -641,8 +639,8 @@ class MobileManipulator:
         # Compute each wheel’s linear speed (m/s) from its angular speed.
         wheel_linear_speeds = wheel_radps * wheel_radius
 
-        # Define the wheel mounting angles with a -90° offset.
-        angles = np.radians(np.array([240, 120, 0]) - 90)
+        # Define the wheel mounting angles (defined from y axis cw)
+        angles = np.radians(np.array([300, 180, 60]))
         m = np.array([[np.cos(a), np.sin(a), base_radius] for a in angles])
 
         # Solve the inverse kinematics: body_velocity = M⁻¹ · wheel_linear_speeds.

lerobot/common/utils/utils.py

@@ -225,13 +225,3 @@ def is_valid_numpy_dtype_string(dtype_str: str) -> bool:
     except TypeError:
         # If a TypeError is raised, the string is not a valid dtype
         return False
-
-
-def get_elapsed_time_in_days_hours_minutes_seconds(elapsed_time_s: float):
-    days = int(elapsed_time_s // (24 * 3600))
-    elapsed_time_s %= 24 * 3600
-    hours = int(elapsed_time_s // 3600)
-    elapsed_time_s %= 3600
-    minutes = int(elapsed_time_s // 60)
-    seconds = elapsed_time_s % 60
-    return days, hours, minutes, seconds

lerobot/configs/policies.py

@@ -47,6 +47,15 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
     def type(self) -> str:
         return self.get_choice_name(self.__class__)
 
+    # TODO(Steven): Find a better way to do deal with this
+    @classmethod
+    def get_type_str(cls) -> str:
+        """Get the policy type identifier for this configuration class."""
+        class_name = cls.__name__.lower()
+        if class_name.endswith("config"):
+            return class_name[:-6]  # Remove 'config' suffix
+        return class_name
+
     @abc.abstractproperty
     def observation_delta_indices(self) -> list | None:
         raise NotImplementedError

lerobot/scripts/visualize_dataset_html.py

@@ -158,7 +158,7 @@ def run_server(
             if major_version < 2:
                 return "Make sure to convert your LeRobotDataset to v2 & above."
 
-        episode_data_csv_str, columns = get_episode_data(dataset, episode_id)
+        episode_data_csv_str, columns, ignored_columns = get_episode_data(dataset, episode_id)
         dataset_info = {
             "repo_id": f"{dataset_namespace}/{dataset_name}",
             "num_samples": dataset.num_frames
@@ -218,6 +218,7 @@ def run_server(
             videos_info=videos_info,
             episode_data_csv_str=episode_data_csv_str,
             columns=columns,
+            ignored_columns=ignored_columns,
         )
 
     app.run(host=host, port=port)
@@ -236,6 +237,14 @@ def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index)
     selected_columns = [col for col, ft in dataset.features.items() if ft["dtype"] == "float32"]
     selected_columns.remove("timestamp")
 
+    ignored_columns = []
+    for column_name in selected_columns:
+        shape = dataset.features[column_name]["shape"]
+        shape_dim = len(shape)
+        if shape_dim > 1:
+            selected_columns.remove(column_name)
+            ignored_columns.append(column_name)
+
     # init header of csv with state and action names
     header = ["timestamp"]
 
@@ -291,7 +300,7 @@ def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index)
     csv_writer.writerows(rows)
     csv_string = csv_buffer.getvalue()
 
-    return csv_string, columns
+    return csv_string, columns, ignored_columns
 
 
 def get_episode_video_paths(dataset: LeRobotDataset, ep_index: int) -> list[str]:

lerobot/templates/visualize_dataset_template.html

@@ -224,49 +224,58 @@
                 </p>
             </div>
 
-            <table class="text-sm border-collapse border border-slate-700" x-show="currentFrameData">
-                <thead>
-                    <tr>
-                        <th></th>
-                        <template x-for="(_, colIndex) in Array.from({length: columns.length}, (_, index) => index)">
-                            <th class="border border-slate-700">
-                                <div class="flex gap-x-2 justify-between px-2">
-                                    <input type="checkbox" :checked="isColumnChecked(colIndex)"
-                                        @change="toggleColumn(colIndex)">
-                                    <p x-text="`${columns[colIndex].key}`"></p>
-                                </div>
-                            </th>
-                        </template>
-                    </tr>
-                </thead>
-                <tbody>
-                    <template x-for="(row, rowIndex) in rows">
-                        <tr class="odd:bg-gray-800 even:bg-gray-900">
-                            <td class="border border-slate-700">
-                                <div class="flex gap-x-2 max-w-64 font-semibold px-1 break-all">
-                                    <input type="checkbox" :checked="isRowChecked(rowIndex)"
-                                        @change="toggleRow(rowIndex)">
-                                </div>
-                            </td>
-                            <template x-for="(cell, colIndex) in row">
-                                <td x-show="cell" class="border border-slate-700">
-                                    <div class="flex gap-x-2 justify-between px-2" :class="{ 'hidden': cell.isNull }">
-                                        <div class="flex gap-x-2">
-                                            <input type="checkbox" x-model="cell.checked" @change="updateTableValues()">
-                                            <span x-text="`${!cell.isNull ? cell.label : null}`"></span>
-                                        </div>
-                                        <span class="w-14 text-right" x-text="`${!cell.isNull ? (typeof cell.value === 'number' ? cell.value.toFixed(2) : cell.value) : null}`"
-                                            :style="`color: ${cell.color}`"></span>
+            <div>
+                <table class="text-sm border-collapse border border-slate-700" x-show="currentFrameData">
+                    <thead>
+                        <tr>
+                            <th></th>
+                            <template x-for="(_, colIndex) in Array.from({length: columns.length}, (_, index) => index)">
+                                <th class="border border-slate-700">
+                                    <div class="flex gap-x-2 justify-between px-2">
+                                        <input type="checkbox" :checked="isColumnChecked(colIndex)"
+                                            @change="toggleColumn(colIndex)">
+                                        <p x-text="`${columns[colIndex].key}`"></p>
                                     </div>
-                                </td>
+                                </th>
                             </template>
                         </tr>
-                    </template>
-                </tbody>
-            </table>
+                    </thead>
+                    <tbody>
+                        <template x-for="(row, rowIndex) in rows">
+                            <tr class="odd:bg-gray-800 even:bg-gray-900">
+                                <td class="border border-slate-700">
+                                    <div class="flex gap-x-2 max-w-64 font-semibold px-1 break-all">
+                                        <input type="checkbox" :checked="isRowChecked(rowIndex)"
+                                            @change="toggleRow(rowIndex)">
+                                    </div>
+                                </td>
+                                <template x-for="(cell, colIndex) in row">
+                                    <td x-show="cell" class="border border-slate-700">
+                                        <div class="flex gap-x-2 justify-between px-2" :class="{ 'hidden': cell.isNull }">
+                                            <div class="flex gap-x-2">
+                                                <input type="checkbox" x-model="cell.checked" @change="updateTableValues()">
+                                                <span x-text="`${!cell.isNull ? cell.label : null}`"></span>
+                                            </div>
+                                            <span class="w-14 text-right" x-text="`${!cell.isNull ? (typeof cell.value === 'number' ? cell.value.toFixed(2) : cell.value) : null}`"
+                                                :style="`color: ${cell.color}`"></span>
+                                        </div>
+                                    </td>
+                                </template>
+                            </tr>
+                        </template>
+                    </tbody>
+                </table>
 
-            <div id="labels" class="hidden">
+                <div id="labels" class="hidden">
+                </div>
+
+                {% if ignored_columns|length > 0 %}
+                <div class="m-2 text-orange-700 max-w-96">
+                    Columns {{ ignored_columns }} are NOT shown since the visualizer currently does not support 2D or 3D data.
+                </div>
+                {% endif %}
             </div>
+
         </div>
     </div>
 

tests/conftest.py

@@ -36,51 +36,27 @@ def pytest_collection_finish():
     print(f"\nTesting with {DEVICE=}")
 
 
-@pytest.fixture
-def is_robot_available(robot_type):
-    if robot_type not in available_robots:
+def _check_component_availability(component_type, available_components, make_component):
+    """Generic helper to check if a hardware component is available"""
+    if component_type not in available_components:
         raise ValueError(
-            f"The robot type '{robot_type}' is not valid. Expected one of these '{available_robots}"
+            f"The {component_type} type is not valid. Expected one of these '{available_components}'"
         )
 
     try:
-        robot = make_robot(robot_type)
-        robot.connect()
-        del robot
+        component = make_component(component_type)
+        component.connect()
+        del component
         return True
 
     except Exception as e:
-        print(f"\nA {robot_type} robot is not available.")
+        print(f"\nA {component_type} is not available.")
 
         if isinstance(e, ModuleNotFoundError):
             print(f"\nInstall module '{e.name}'")
         elif isinstance(e, SerialException):
-            print("\nNo physical motors bus detected.")
-        else:
-            traceback.print_exc()
-
-        return False
-
-
-@pytest.fixture
-def is_camera_available(camera_type):
-    if camera_type not in available_cameras:
-        raise ValueError(
-            f"The camera type '{camera_type}' is not valid. Expected one of these '{available_cameras}"
-        )
-
-    try:
-        camera = make_camera(camera_type)
-        camera.connect()
-        del camera
-        return True
-
-    except Exception as e:
-        print(f"\nA {camera_type} camera is not available.")
-
-        if isinstance(e, ModuleNotFoundError):
-            print(f"\nInstall module '{e.name}'")
-        elif isinstance(e, ValueError) and "camera_index" in e.args[0]:
+            print("\nNo physical device detected.")
+        elif isinstance(e, ValueError) and "camera_index" in str(e):
             print("\nNo physical camera detected.")
         else:
             traceback.print_exc()
@@ -88,30 +64,19 @@ def is_camera_available(camera_type):
         return False
 
 
+@pytest.fixture
+def is_robot_available(robot_type):
+    return _check_component_availability(robot_type, available_robots, make_robot)
+
+
+@pytest.fixture
+def is_camera_available(camera_type):
+    return _check_component_availability(camera_type, available_cameras, make_camera)
+
+
 @pytest.fixture
 def is_motor_available(motor_type):
-    if motor_type not in available_motors:
-        raise ValueError(
-            f"The motor type '{motor_type}' is not valid. Expected one of these '{available_motors}"
-        )
-
-    try:
-        motors_bus = make_motors_bus(motor_type)
-        motors_bus.connect()
-        del motors_bus
-        return True
-
-    except Exception as e:
-        print(f"\nA {motor_type} motor is not available.")
-
-        if isinstance(e, ModuleNotFoundError):
-            print(f"\nInstall module '{e.name}'")
-        elif isinstance(e, SerialException):
-            print("\nNo physical motors bus detected.")
-        else:
-            traceback.print_exc()
-
-        return False
+    return _check_component_availability(motor_type, available_motors, make_motors_bus)
 
 
 @pytest.fixture

tests/test_aggregate_datasets.py

@@ -1,19 +0,0 @@
-from lerobot.common.datasets.aggregate import aggregate_datasets
-from tests.fixtures.constants import DUMMY_REPO_ID
-
-
-def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
-    dataset_0 = lerobot_dataset_factory(
-        root=tmp_path / "test_0",
-        repo_id=DUMMY_REPO_ID + "_0",
-        total_episodes=10,
-        total_frames=400,
-    )
-    dataset_1 = lerobot_dataset_factory(
-        root=tmp_path / "test_1",
-        repo_id=DUMMY_REPO_ID + "_1",
-        total_episodes=10,
-        total_frames=400,
-    )
-
-    dataset_2 = aggregate_datasets([dataset_0, dataset_1])