- We just added a new tutorial on how to build a more affordable robot, at the price of $110 per arm!
-Teach it new skills by showing it a few moves with just a laptop.
-Then watch your homemade robot act autonomously 🤯
-Follow the link to the full tutorial for SO-100.
+
+
+ Meet the SO-100 – Just $110 per arm!
+Train it in minutes with a few simple moves on your laptop.
+Then sit back and watch your creation act autonomously! 🤯
+ ++ Get the full SO-100 tutorial here.
+ +Want to take it to the next level? Make your SO-100 mobile by building LeKiwi!
+Check out the LeKiwi tutorial and bring your robot to life on wheels.
+ +
+
+**Step 3: Install in Base**
+- Place the first motor into the base.
+
+ 
+
+**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).
+
+ 
+
+**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.
+
+ 
+
+
+**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.
+
+ 
+
+**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.
+
+ 
+ 
+ 
+ 
+
+**Step 14: Attach Upper Arm**
+- Attach the upper arm with 4 screws on each side.
+
+ 
+
+---
+
+### 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.
+
+ 
+
+**Step 17: Attach Forearm**
+- Connect the forearm to motor 3 using 4 screws on each side.
+
+ 
+
+---
+
+### 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.
+
+ 
+ 
+ 
+
+**Step 20: Secure Motor 4 & Attach Horn**
+- Fasten motor 4 with 4 screws and attach its motor horns, use for one a horn screw.
+
+ 
+
+---
+
+### Wrist Assembly
+
+**Step 21: Install Motor 5**
+- Insert motor 5 into the wrist holder and secure it with 2 front screws.
+
+ 
+
+**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.
+
+ 
+
+**Step 23: Attach Wrist Horn**
+- Install only one motor horn on the wrist motor and secure it with a horn screw.
+
+ 
+
+---
+
+### Follower Configuration
+
+**Step 24: Attach Gripper**
+- Attach the gripper to motor 5.
+
+ 
+
+**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.
+
+ 
+
+**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.
+
+ 
+
+**Step 27: Mount Controller**
+- Attach the motor controller on the back.
+
+ 
+ 
+ 
+
+**Step 25: Attach Handle**
+- Attach the handle to motor 5 using 4 screws.
+
+ 
+
+**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.
+
+ 
+
+**Step 27: Attach Trigger**
+- Attach the follower trigger with 4 screws.
+
+ 
+
+**Step 28: Mount Controller**
+- Attach the motor controller on the back.
+
+ 
| 
| 
|
Make sure both arms are connected and run this script to launch manual calibration:
@@ -271,8 +470,8 @@ python lerobot/scripts/control_robot.py \
#### b. Manual calibration of leader arm
Follow step 6 of the [assembly video](https://youtu.be/FioA2oeFZ5I?t=724) which illustrates the manual calibration. You will need to move the leader arm to these positions sequentially:
-| 1. Zero position | 2. Rotated position | 3. Rest position |
-|---|---|---|
+| 1. Zero position | 2. Rotated position | 3. Rest position |
+| ------------------------------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------ |
| 
| 
| 
|
Run this script to launch manual calibration:
@@ -335,7 +534,7 @@ python lerobot/scripts/control_robot.py \
--control.push_to_hub=true
```
-Note: You can resume recording by adding `--control.resume=true`. Also if you didn't push your dataset yet, add `--control.local_files_only=true`.
+Note: You can resume recording by adding `--control.resume=true`.
## H. Visualize a dataset
@@ -344,7 +543,7 @@ If you uploaded your dataset to the hub with `--control.push_to_hub=true`, you c
echo ${HF_USER}/so100_test
```
-If you didn't upload with `--control.push_to_hub=false`, you can also visualize it locally with:
+If you didn't upload with `--control.push_to_hub=false`, you can also visualize it locally with (a window can be opened in the browser `http://127.0.0.1:9090` with the visualization tool):
```bash
python lerobot/scripts/visualize_dataset_html.py \
--repo-id ${HF_USER}/so100_test \
@@ -363,8 +562,6 @@ python lerobot/scripts/control_robot.py \
--control.episode=0
```
-Note: If you didn't push your dataset yet, add `--control.local_files_only=true`.
-
## J. Train a policy
To train a policy to control your robot, use the [`python lerobot/scripts/train.py`](../lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
@@ -374,20 +571,25 @@ python lerobot/scripts/train.py \
--policy.type=act \
--output_dir=outputs/train/act_so100_test \
--job_name=act_so100_test \
- --device=cuda \
+ --policy.device=cuda \
--wandb.enable=true
```
-Note: If you didn't push your dataset yet, add `--control.local_files_only=true`.
-
Let's explain it:
1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/so100_test`.
2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../lerobot/common/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor sates, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
-4. We provided `device=cuda` since we are training on a Nvidia GPU, but you could use `device=mps` to train on Apple silicon.
+4. We provided `policy.device=cuda` since we are training on a Nvidia GPU, but you could use `policy.device=mps` to train on Apple silicon.
5. We provided `wandb.enable=true` to use [Weights and Biases](https://docs.wandb.ai/quickstart) for visualizing training plots. This is optional but if you use it, make sure you are logged in by running `wandb login`.
Training should take several hours. You will find checkpoints in `outputs/train/act_so100_test/checkpoints`.
+To resume training from a checkpoint, below is an example command to resume from `last` checkpoint of the `act_so100_test` policy:
+```bash
+python lerobot/scripts/train.py \
+ --config_path=outputs/train/act_so100_test/checkpoints/last/pretrained_model/train_config.json \
+ --resume=true
+```
+
## K. Evaluate your policy
You can use the `record` function from [`lerobot/scripts/control_robot.py`](../lerobot/scripts/control_robot.py) but with a policy checkpoint as input. For instance, run this command to record 10 evaluation episodes:
@@ -416,4 +618,4 @@ As you can see, it's almost the same command as previously used to record your t
Follow this [previous tutorial](https://github.com/huggingface/lerobot/blob/main/examples/7_get_started_with_real_robot.md#4-train-a-policy-on-your-data) for a more in-depth tutorial on controlling real robots with LeRobot.
> [!TIP]
-> If you have any questions or need help, please reach out on Discord in the channel [`#so100-arm`](https://discord.com/channels/1216765309076115607/1237741463832363039).
+> If you have any questions or need help, please reach out on [Discord](https://discord.com/invite/s3KuuzsPFb) in the channel [`#so100-arm`](https://discord.com/channels/1216765309076115607/1237741463832363039).
diff --git a/examples/11_use_lekiwi.md b/examples/11_use_lekiwi.md
new file mode 100644
index 00000000..dc310af2
--- /dev/null
+++ b/examples/11_use_lekiwi.md
@@ -0,0 +1,585 @@
+# Using the [LeKiwi](https://github.com/SIGRobotics-UIUC/LeKiwi) Robot with LeRobot
+
+## Table of Contents
+
+ - [A. Source the parts](#a-source-the-parts)
+ - [B. Install software Pi](#b-install-software-on-pi)
+ - [C. Setup LeRobot laptop/pc](#c-install-lerobot-on-laptop)
+ - [D. Assemble the arms](#d-assembly)
+ - [E. Calibrate](#e-calibration)
+ - [F. Teleoperate](#f-teleoperate)
+ - [G. Record a dataset](#g-record-a-dataset)
+ - [H. Visualize a dataset](#h-visualize-a-dataset)
+ - [I. Replay an episode](#i-replay-an-episode)
+ - [J. Train a policy](#j-train-a-policy)
+ - [K. Evaluate your policy](#k-evaluate-your-policy)
+
+> [!TIP]
+> If you have any questions or need help, please reach out on [Discord](https://discord.com/invite/s3KuuzsPFb) in the channel [`#mobile-so-100-arm`](https://discord.com/channels/1216765309076115607/1318390825528332371).
+
+## A. Source the parts
+
+Follow this [README](https://github.com/SIGRobotics-UIUC/LeKiwi). It contains the bill of materials, with a link to source the parts, as well as the instructions to 3D print the parts, and advice if it's your first time printing or if you don't own a 3D printer.
+
+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.
+
+### Install OS
+For setting up the Raspberry Pi and its SD-card see: [Setup PI](https://www.raspberrypi.com/documentation/computers/getting-started.html). Here is explained how to download the [Imager](https://www.raspberrypi.com/software/) to install Raspberry Pi OS or Ubuntu.
+
+### Setup SSH
+After setting up your Pi, you should enable and setup [SSH](https://www.raspberrypi.com/news/coding-on-raspberry-pi-remotely-with-visual-studio-code/) (Secure Shell Protocol) so you can login into the Pi from your laptop without requiring a screen, keyboard and mouse in the Pi. A great tutorial on how to do this can be found [here](https://www.raspberrypi.com/documentation/computers/remote-access.html#ssh). Logging into your Pi can be done in your Command Prompt (cmd) or if you use VSCode you can use [this](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-ssh) extension.
+
+### Install LeRobot
+
+On your Raspberry Pi:
+
+#### 1. [Install Miniconda](https://docs.anaconda.com/miniconda/install/#quick-command-line-install):
+
+#### 2. Restart shell
+Copy paste in your shell: `source ~/.bashrc` or for Mac: `source ~/.bash_profile` or `source ~/.zshrc` if you're using zshell
+
+#### 3. Create and activate a fresh conda environment for lerobot
+
+
+
+### Assemble arms
+[Assemble arms instruction](https://github.com/huggingface/lerobot/blob/main/examples/10_use_so100.md#d-assemble-the-arms)
+
+## Mobile base (LeKiwi)
+[Assemble LeKiwi](https://github.com/SIGRobotics-UIUC/LeKiwi)
+
+### Update config
+Both config files on the LeKiwi LeRobot and on the laptop should be the same. First we should find the Ip address of the Raspberry Pi of the mobile manipulator. This is the same Ip address used in SSH. We also need the usb port of the control board of the leader arm on the laptop and the port of the control board on LeKiwi. We can find these ports with the following script.
+
+#### a. Run the script to find port
+
+
| 
| 
|
+
+Make sure the arm is connected to the Raspberry Pi and run this script (on the Raspberry Pi) to launch manual calibration:
+```bash
+python lerobot/scripts/control_robot.py \
+ --robot.type=lekiwi \
+ --robot.cameras='{}' \
+ --control.type=calibrate \
+ --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:
+
+| 1. Zero position | 2. Rotated position | 3. Rest position |
+| ------------------------------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| 
| 
| 
|
Make sure both arms are connected and run this script to launch manual calibration:
@@ -192,8 +192,8 @@ python lerobot/scripts/control_robot.py \
**Manual calibration of leader arm**
Follow step 6 of the [assembly video](https://www.youtube.com/watch?v=DA91NJOtMic) which illustrates the manual calibration. You will need to move the leader arm to these positions sequentially:
-| 1. Zero position | 2. Rotated position | 3. Rest position |
-|---|---|---|
+| 1. Zero position | 2. Rotated position | 3. Rest position |
+| ------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------- |
| 
| 
| 
|
Run this script to launch manual calibration:
@@ -256,7 +256,7 @@ python lerobot/scripts/control_robot.py \
--control.push_to_hub=true
```
-Note: You can resume recording by adding `--control.resume=true`. Also if you didn't push your dataset yet, add `--control.local_files_only=true`.
+Note: You can resume recording by adding `--control.resume=true`.
## Visualize a dataset
@@ -284,8 +284,6 @@ python lerobot/scripts/control_robot.py \
--control.episode=0
```
-Note: If you didn't push your dataset yet, add `--control.local_files_only=true`.
-
## Train a policy
To train a policy to control your robot, use the [`python lerobot/scripts/train.py`](../lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
@@ -295,16 +293,14 @@ python lerobot/scripts/train.py \
--policy.type=act \
--output_dir=outputs/train/act_moss_test \
--job_name=act_moss_test \
- --device=cuda \
+ --policy.device=cuda \
--wandb.enable=true
```
-Note: If you didn't push your dataset yet, add `--control.local_files_only=true`.
-
Let's explain it:
1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/moss_test`.
2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../lerobot/common/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor sates, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
-4. We provided `device=cuda` since we are training on a Nvidia GPU, but you could use `device=mps` to train on Apple silicon.
+4. We provided `policy.device=cuda` since we are training on a Nvidia GPU, but you could use `policy.device=mps` to train on Apple silicon.
5. We provided `wandb.enable=true` to use [Weights and Biases](https://docs.wandb.ai/quickstart) for visualizing training plots. This is optional but if you use it, make sure you are logged in by running `wandb login`.
Training should take several hours. You will find checkpoints in `outputs/train/act_moss_test/checkpoints`.
diff --git a/examples/1_load_lerobot_dataset.py b/examples/1_load_lerobot_dataset.py
index 96c104b6..c374a375 100644
--- a/examples/1_load_lerobot_dataset.py
+++ b/examples/1_load_lerobot_dataset.py
@@ -1,3 +1,17 @@
+# 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.
+
"""
This script demonstrates the use of `LeRobotDataset` class for handling and processing robotic datasets from Hugging Face.
It illustrates how to load datasets, manipulate them, and apply transformations suitable for machine learning tasks in PyTorch.
diff --git a/examples/2_evaluate_pretrained_policy.py b/examples/2_evaluate_pretrained_policy.py
index 0a7b8deb..edbbad38 100644
--- a/examples/2_evaluate_pretrained_policy.py
+++ b/examples/2_evaluate_pretrained_policy.py
@@ -1,3 +1,17 @@
+# 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.
+
"""
This scripts demonstrates how to evaluate a pretrained policy from the HuggingFace Hub or from your local
training outputs directory. In the latter case, you might want to run examples/3_train_policy.py first.
@@ -30,7 +44,7 @@ pretrained_policy_path = "lerobot/diffusion_pusht"
# OR a path to a local outputs/train folder.
# pretrained_policy_path = Path("outputs/train/example_pusht_diffusion")
-policy = DiffusionPolicy.from_pretrained(pretrained_policy_path, map_location=device)
+policy = DiffusionPolicy.from_pretrained(pretrained_policy_path)
# Initialize evaluation environment to render two observation types:
# an image of the scene and state/position of the agent. The environment
diff --git a/examples/3_train_policy.py b/examples/3_train_policy.py
index 635c7293..6c3af54e 100644
--- a/examples/3_train_policy.py
+++ b/examples/3_train_policy.py
@@ -1,3 +1,17 @@
+# 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.
+
"""This scripts demonstrates how to train Diffusion Policy on the PushT environment.
Once you have trained a model with this script, you can try to evaluate it on
@@ -85,9 +99,8 @@ 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()}
- output_dict = policy.forward(batch)
- loss = output_dict["loss"]
+ 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()
optimizer.zero_grad()
diff --git a/examples/4_train_policy_with_script.md b/examples/4_train_policy_with_script.md
index 9d57d424..b23d2271 100644
--- a/examples/4_train_policy_with_script.md
+++ b/examples/4_train_policy_with_script.md
@@ -1,5 +1,5 @@
This tutorial will explain the training script, how to use it, and particularly how to configure everything needed for the training run.
-> **Note:** The following assume you're running these commands on a machine equipped with a cuda GPU. If you don't have one (or if you're using a Mac), you can add `--device=cpu` (`--device=mps` respectively). However, be advised that the code executes much slower on cpu.
+> **Note:** The following assume you're running these commands on a machine equipped with a cuda GPU. If you don't have one (or if you're using a Mac), you can add `--policy.device=cpu` (`--policy.device=mps` respectively). However, be advised that the code executes much slower on cpu.
## The training script
@@ -161,13 +161,13 @@ python lerobot/scripts/train.py \
```
You should see from the logging that your training picks up from where it left off.
-Another reason for which you might want to resume a run is simply to extend training and add more training steps. The number of training steps is set by the option `--offline.steps`, which is 100 000 by default.
+Another reason for which you might want to resume a run is simply to extend training and add more training steps. The number of training steps is set by the option `--steps`, which is 100 000 by default.
You could double the number of steps of the previous run with:
```bash
python lerobot/scripts/train.py \
--config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
--resume=true \
- --offline.steps=200000
+ --steps=200000
```
## Outputs of a run
@@ -175,12 +175,16 @@ In the output directory, there will be a folder called `checkpoints` with the fo
```bash
outputs/train/run_resumption/checkpoints
├── 000100 # checkpoint_dir for training step 100
-│ ├── pretrained_model
-│ │ ├── config.json # pretrained policy config
-│ │ ├── model.safetensors # model weights
-│ │ ├── train_config.json # train config
-│ │ └── README.md # model card
-│ └── training_state.pth # optimizer/scheduler/rng state and training step
+│ ├── pretrained_model/
+│ │ ├── config.json # policy config
+│ │ ├── model.safetensors # policy weights
+│ │ └── train_config.json # train config
+│ └── training_state/
+│ ├── optimizer_param_groups.json # optimizer param groups
+│ ├── optimizer_state.safetensors # optimizer state
+│ ├── rng_state.safetensors # rng states
+│ ├── scheduler_state.json # scheduler state
+│ └── training_step.json # training step
├── 000200
└── last -> 000200 # symlink to the last available checkpoint
```
@@ -250,7 +254,7 @@ python lerobot/scripts/train.py \
python lerobot/scripts/train.py \
--config_path=checkpoint/pretrained_model/ \
--resume=true \
- --offline.steps=200000 # <- you can change some training parameters
+ --steps=200000 # <- you can change some training parameters
```
#### Fine-tuning
diff --git a/examples/7_get_started_with_real_robot.md b/examples/7_get_started_with_real_robot.md
index 1f4c7aaa..52ec673e 100644
--- a/examples/7_get_started_with_real_robot.md
+++ b/examples/7_get_started_with_real_robot.md
@@ -36,9 +36,14 @@ Using `pip`:
pip install -e ".[dynamixel]"
```
-Or using `poetry`:
+Using `poetry`:
```bash
-poetry install --sync --extras "dynamixel"
+poetry sync --extras "dynamixel"
+```
+
+Using `uv`:
+```bash
+uv sync --extra "dynamixel"
```
/!\ For Linux only, ffmpeg and opencv requires conda install for now. Run this exact sequence of commands:
@@ -287,6 +292,11 @@ Steps:
- Scan for devices. All 12 motors should appear.
- Select the motors one by one and move the arm. Check that the graphical indicator near the top right shows the movement.
+** There is a common issue with the Dynamixel XL430-W250 motors where the motors become undiscoverable after upgrading their firmware from Mac and Windows Dynamixel Wizard2 applications. When this occurs, it is required to do a firmware recovery (Select `DYNAMIXEL Firmware Recovery` and follow the prompts). There are two known workarounds to conduct this firmware reset:
+ 1) Install the Dynamixel Wizard on a linux machine and complete the firmware recovery
+ 2) Use the Dynamixel U2D2 in order to perform the reset with Windows or Mac. This U2D2 can be purchased [here](https://www.robotis.us/u2d2/).
+ For either solution, open DYNAMIXEL Wizard 2.0 and select the appropriate port. You will likely be unable to see the motor in the GUI at this time. Select `Firmware Recovery`, carefully choose the correct model, and wait for the process to complete. Finally, re-scan to confirm the firmware recovery was successful.
+
**Read and Write with DynamixelMotorsBus**
To get familiar with how `DynamixelMotorsBus` communicates with the motors, you can start by reading data from them. Copy past this code in the same interactive python session:
@@ -381,19 +391,19 @@ When you connect your robot for the first time, the [`ManipulatorRobot`](../lero
Here are the positions you'll move the follower arm to:
-| 1. Zero position | 2. Rotated position | 3. Rest position |
-|---|---|---|
+| 1. Zero position | 2. Rotated position | 3. Rest position |
+| ----------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| 
| 
| 
|
And here are the corresponding positions for the leader arm:
-| 1. Zero position | 2. Rotated position | 3. Rest position |
-|---|---|---|
+| 1. Zero position | 2. Rotated position | 3. Rest position |
+| ----------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------- |
| 
| 
| 
|
You can watch a [video tutorial of the calibration procedure](https://youtu.be/8drnU9uRY24) for more details.
-During calibration, we count the number of full 360-degree rotations your motors have made since they were first used. That's why we ask yo to move to this arbitrary "zero" position. We don't actually "set" the zero position, so you don't need to be accurate. After calculating these "offsets" to shift the motor values around 0, we need to assess the rotation direction of each motor, which might differ. That's why we ask you to rotate all motors to roughly 90 degrees, to mesure if the values changed negatively or positively.
+During calibration, we count the number of full 360-degree rotations your motors have made since they were first used. That's why we ask yo to move to this arbitrary "zero" position. We don't actually "set" the zero position, so you don't need to be accurate. After calculating these "offsets" to shift the motor values around 0, we need to assess the rotation direction of each motor, which might differ. That's why we ask you to rotate all motors to roughly 90 degrees, to measure if the values changed negatively or positively.
Finally, the rest position ensures that the follower and leader arms are roughly aligned after calibration, preventing sudden movements that could damage the motors when starting teleoperation.
@@ -621,7 +631,7 @@ Finally, run this code to instantiate and connectyour camera:
from lerobot.common.robot_devices.cameras.configs import OpenCVCameraConfig
from lerobot.common.robot_devices.cameras.opencv import OpenCVCamera
-camera_config = OpenCVCameraConfig(camera_index=0)
+config = OpenCVCameraConfig(camera_index=0)
camera = OpenCVCamera(config)
camera.connect()
color_image = camera.read()
@@ -658,18 +668,20 @@ camera.disconnect()
**Instantiate your robot with cameras**
-Additionaly, you can set up your robot to work with your cameras.
+Additionally, you can set up your robot to work with your cameras.
Modify the following Python code with the appropriate camera names and configurations:
```python
robot = ManipulatorRobot(
- leader_arms={"main": leader_arm},
- follower_arms={"main": follower_arm},
- calibration_dir=".cache/calibration/koch",
- cameras={
- "laptop": OpenCVCameraConfig(0, fps=30, width=640, height=480),
- "phone": OpenCVCameraConfig(1, fps=30, width=640, height=480),
- },
+ KochRobotConfig(
+ leader_arms={"main": leader_arm},
+ follower_arms={"main": follower_arm},
+ calibration_dir=".cache/calibration/koch",
+ cameras={
+ "laptop": OpenCVCameraConfig(0, fps=30, width=640, height=480),
+ "phone": OpenCVCameraConfig(1, fps=30, width=640, height=480),
+ },
+ )
)
robot.connect()
```
@@ -706,7 +718,7 @@ python lerobot/scripts/control_robot.py \
You will see a lot of lines appearing like this one:
```
-INFO 2024-08-10 11:15:03 ol_robot.py:209 dt: 5.12 (195.1hz) dtRlead: 4.93 (203.0hz) dtRfoll: 0.19 (5239.0hz)
+INFO 2024-08-10 11:15:03 ol_robot.py:209 dt: 5.12 (195.1hz) dtRlead: 4.93 (203.0hz) dtWfoll: 0.19 (5239.0hz)
```
It contains
@@ -763,7 +775,7 @@ You can use the `record` function from [`lerobot/scripts/control_robot.py`](../l
1. Frames from cameras are saved on disk in threads, and encoded into videos at the end of each episode recording.
2. Video streams from cameras are displayed in window so that you can verify them.
3. Data is stored with [`LeRobotDataset`](../lerobot/common/datasets/lerobot_dataset.py) format which is pushed to your Hugging Face page (unless `--control.push_to_hub=false` is provided).
-4. Checkpoints are done during recording, so if any issue occurs, you can resume recording by re-running the same command again with `--control.resume=true`. You might need to add `--control.local_files_only=true` if your dataset was not uploaded to hugging face hub. Also you will need to manually delete the dataset directory to start recording from scratch.
+4. Checkpoints are done during recording, so if any issue occurs, you can resume recording by re-running the same command again with `--control.resume=true`. You will need to manually delete the dataset directory if you want to start recording from scratch.
5. Set the flow of data recording using command line arguments:
- `--control.warmup_time_s=10` defines the number of seconds before starting data collection. It allows the robot devices to warmup and synchronize (10 seconds by default).
- `--control.episode_time_s=60` defines the number of seconds for data recording for each episode (60 seconds by default).
@@ -818,8 +830,8 @@ It contains:
- `dtRlead: 5.06 (197.5hz)` which is the delta time of reading the present position of the leader arm.
- `dtWfoll: 0.25 (3963.7hz)` which is the delta time of writing the goal position on the follower arm ; writing is asynchronous so it takes less time than reading.
- `dtRfoll: 6.22 (160.7hz)` which is the delta time of reading the present position on the follower arm.
-- `dtRlaptop:32.57 (30.7hz) ` which is the delta time of capturing an image from the laptop camera in the thread running asynchrously.
-- `dtRphone:33.84 (29.5hz)` which is the delta time of capturing an image from the phone camera in the thread running asynchrously.
+- `dtRlaptop:32.57 (30.7hz) ` which is the delta time of capturing an image from the laptop camera in the thread running asynchronously.
+- `dtRphone:33.84 (29.5hz)` which is the delta time of capturing an image from the phone camera in the thread running asynchronously.
Troubleshooting:
- On Linux, if you encounter a hanging issue when using cameras, uninstall opencv and re-install it with conda:
@@ -839,7 +851,7 @@ At the end of data recording, your dataset will be uploaded on your Hugging Face
echo https://huggingface.co/datasets/${HF_USER}/koch_test
```
-### b. Advices for recording dataset
+### b. Advice for recording dataset
Once you're comfortable with data recording, it's time to create a larger dataset for training. A good starting task is grasping an object at different locations and placing it in a bin. We suggest recording at least 50 episodes, with 10 episodes per location. Keep the cameras fixed and maintain consistent grasping behavior throughout the recordings.
@@ -878,8 +890,6 @@ python lerobot/scripts/control_robot.py \
--control.episode=0
```
-Note: You might need to add `--control.local_files_only=true` if your dataset was not uploaded to hugging face hub.
-
Your robot should replicate movements similar to those you recorded. For example, check out [this video](https://x.com/RemiCadene/status/1793654950905680090) where we use `replay` on a Aloha robot from [Trossen Robotics](https://www.trossenrobotics.com).
## 4. Train a policy on your data
@@ -893,16 +903,14 @@ python lerobot/scripts/train.py \
--policy.type=act \
--output_dir=outputs/train/act_koch_test \
--job_name=act_koch_test \
- --device=cuda \
+ --policy.device=cuda \
--wandb.enable=true
```
-Note: You might need to add `--dataset.local_files_only=true` if your dataset was not uploaded to hugging face hub.
-
Let's explain it:
1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/koch_test`.
2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../lerobot/common/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor sates, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
-4. We provided `device=cuda` since we are training on a Nvidia GPU, but you could use `device=mps` to train on Apple silicon.
+4. We provided `policy.device=cuda` since we are training on a Nvidia GPU, but you could use `policy.device=mps` to train on Apple silicon.
5. We provided `wandb.enable=true` to use [Weights and Biases](https://docs.wandb.ai/quickstart) for visualizing training plots. This is optional but if you use it, make sure you are logged in by running `wandb login`.
For more information on the `train` script see the previous tutorial: [`examples/4_train_policy_with_script.md`](../examples/4_train_policy_with_script.md)
diff --git a/examples/8_use_stretch.md b/examples/8_use_stretch.md
index 2f8c0ffb..802ea718 100644
--- a/examples/8_use_stretch.md
+++ b/examples/8_use_stretch.md
@@ -98,7 +98,7 @@ python lerobot/scripts/control_robot.py \
```
This is equivalent to running `stretch_robot_home.py`
-> **Note:** If you run any of the LeRobot scripts below and Stretch is not poperly homed, it will automatically home/calibrate first.
+> **Note:** If you run any of the LeRobot scripts below and Stretch is not properly homed, it will automatically home/calibrate first.
**Teleoperate**
Before trying teleoperation, you need activate the gamepad controller by pressing the middle button. For more info, see Stretch's [doc](https://docs.hello-robot.com/0.3/getting_started/hello_robot/#gamepad-teleoperation).
diff --git a/examples/9_use_aloha.md b/examples/9_use_aloha.md
index d74c8b7a..62dee588 100644
--- a/examples/9_use_aloha.md
+++ b/examples/9_use_aloha.md
@@ -2,7 +2,7 @@ This tutorial explains how to use [Aloha and Aloha 2 stationary](https://www.tro
## Setup
-Follow the [documentation from Trossen Robotics](https://docs.trossenrobotics.com/aloha_docs/getting_started/stationary/hardware_setup.html) for setting up the hardware and plugging the 4 arms and 4 cameras to your computer.
+Follow the [documentation from Trossen Robotics](https://docs.trossenrobotics.com/aloha_docs/2.0/getting_started/stationary/hardware_setup.html) for setting up the hardware and plugging the 4 arms and 4 cameras to your computer.
## Install LeRobot
@@ -135,14 +135,14 @@ python lerobot/scripts/train.py \
--policy.type=act \
--output_dir=outputs/train/act_aloha_test \
--job_name=act_aloha_test \
- --device=cuda \
+ --policy.device=cuda \
--wandb.enable=true
```
Let's explain it:
1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/aloha_test`.
2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../lerobot/common/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor sates, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
-4. We provided `device=cuda` since we are training on a Nvidia GPU, but you could use `device=mps` to train on Apple silicon.
+4. We provided `policy.device=cuda` since we are training on a Nvidia GPU, but you could use `policy.device=mps` to train on Apple silicon.
5. We provided `wandb.enable=true` to use [Weights and Biases](https://docs.wandb.ai/quickstart) for visualizing training plots. This is optional but if you use it, make sure you are logged in by running `wandb login`.
For more information on the `train` script see the previous tutorial: [`examples/4_train_policy_with_script.md`](../examples/4_train_policy_with_script.md)
@@ -172,10 +172,10 @@ python lerobot/scripts/control_robot.py \
As you can see, it's almost the same command as previously used to record your training dataset. Two things changed:
1. There is an additional `--control.policy.path` argument which indicates the path to your policy checkpoint with (e.g. `outputs/train/eval_act_aloha_test/checkpoints/last/pretrained_model`). You can also use the model repository if you uploaded a model checkpoint to the hub (e.g. `${HF_USER}/act_aloha_test`).
2. The name of dataset begins by `eval` to reflect that you are running inference (e.g. `${HF_USER}/eval_act_aloha_test`).
-3. We use `--control.num_image_writer_processes=1` instead of the default value (`0`). On our computer, using a dedicated process to write images from the 4 cameras on disk allows to reach constent 30 fps during inference. Feel free to explore different values for `--control.num_image_writer_processes`.
+3. We use `--control.num_image_writer_processes=1` instead of the default value (`0`). On our computer, using a dedicated process to write images from the 4 cameras on disk allows to reach constant 30 fps during inference. Feel free to explore different values for `--control.num_image_writer_processes`.
## More
-Follow this [previous tutorial](https://github.com/huggingface/lerobot/blob/main/examples/7_get_started_with_real_robot.md#4-train-a-policy-on-your-data) for a more in-depth explaination.
+Follow this [previous tutorial](https://github.com/huggingface/lerobot/blob/main/examples/7_get_started_with_real_robot.md#4-train-a-policy-on-your-data) for a more in-depth explanation.
If you have any question or need help, please reach out on Discord in the channel `#aloha-arm`.
diff --git a/examples/advanced/1_add_image_transforms.py b/examples/advanced/1_add_image_transforms.py
index 882710e3..f1460926 100644
--- a/examples/advanced/1_add_image_transforms.py
+++ b/examples/advanced/1_add_image_transforms.py
@@ -1,3 +1,17 @@
+# 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.
+
"""
This script demonstrates how to use torchvision's image transformation with LeRobotDataset for data
augmentation purposes. The transformations are passed to the dataset as an argument upon creation, and
diff --git a/examples/advanced/2_calculate_validation_loss.py b/examples/advanced/2_calculate_validation_loss.py
index 71e76072..47b4dd02 100644
--- a/examples/advanced/2_calculate_validation_loss.py
+++ b/examples/advanced/2_calculate_validation_loss.py
@@ -1,3 +1,17 @@
+# 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.
+
"""This script demonstrates how to slice a dataset and calculate the loss on a subset of the data.
This technique can be useful for debugging and testing purposes, as well as identifying whether a policy
@@ -75,9 +89,9 @@ def main():
n_examples_evaluated = 0
for batch in val_dataloader:
batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
- output_dict = policy.forward(batch)
+ loss, _ = policy.forward(batch)
- loss_cumsum += output_dict["loss"].item()
+ loss_cumsum += loss.item()
n_examples_evaluated += batch["index"].shape[0]
# Calculate the average loss over the validation set.
diff --git a/examples/port_datasets/pusht_zarr.py b/examples/port_datasets/pusht_zarr.py
index 1506f427..ea2e8b60 100644
--- a/examples/port_datasets/pusht_zarr.py
+++ b/examples/port_datasets/pusht_zarr.py
@@ -1,10 +1,25 @@
+# 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 shutil
from pathlib import Path
import numpy as np
-import torch
+from huggingface_hub import HfApi
-from lerobot.common.datasets.lerobot_dataset import LEROBOT_HOME, LeRobotDataset
+from lerobot.common.constants import HF_LEROBOT_HOME
+from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
from lerobot.common.datasets.push_dataset_to_hub._download_raw import download_raw
PUSHT_TASK = "Push the T-shaped blue block onto the T-shaped green target surface."
@@ -89,9 +104,9 @@ def calculate_coverage(zarr_data):
num_frames = len(block_pos)
- coverage = np.zeros((num_frames,))
+ coverage = np.zeros((num_frames,), dtype=np.float32)
# 8 keypoints with 2 coords each
- keypoints = np.zeros((num_frames, 16))
+ keypoints = np.zeros((num_frames, 16), dtype=np.float32)
# Set x, y, theta (in radians)
goal_pos_angle = np.array([256, 256, np.pi / 4])
@@ -117,7 +132,7 @@ def calculate_coverage(zarr_data):
intersection_area = goal_geom.intersection(block_geom).area
goal_area = goal_geom.area
coverage[i] = intersection_area / goal_area
- keypoints[i] = torch.from_numpy(PushTEnv.get_keypoints(block_shapes).flatten())
+ keypoints[i] = PushTEnv.get_keypoints(block_shapes).flatten()
return coverage, keypoints
@@ -134,8 +149,8 @@ def main(raw_dir: Path, repo_id: str, mode: str = "video", push_to_hub: bool = T
if mode not in ["video", "image", "keypoints"]:
raise ValueError(mode)
- if (LEROBOT_HOME / repo_id).exists():
- shutil.rmtree(LEROBOT_HOME / repo_id)
+ if (HF_LEROBOT_HOME / repo_id).exists():
+ shutil.rmtree(HF_LEROBOT_HOME / repo_id)
if not raw_dir.exists():
download_raw(raw_dir, repo_id="lerobot-raw/pusht_raw")
@@ -148,6 +163,10 @@ def main(raw_dir: Path, repo_id: str, mode: str = "video", push_to_hub: bool = T
action = zarr_data["action"][:]
image = zarr_data["img"] # (b, h, w, c)
+ if image.dtype == np.float32 and image.max() == np.float32(255):
+ # HACK: images are loaded as float32 but they actually encode uint8 data
+ image = image.astype(np.uint8)
+
episode_data_index = {
"from": np.concatenate(([0], zarr_data.meta["episode_ends"][:-1])),
"to": zarr_data.meta["episode_ends"],
@@ -175,28 +194,30 @@ def main(raw_dir: Path, repo_id: str, mode: str = "video", push_to_hub: bool = T
for frame_idx in range(num_frames):
i = from_idx + frame_idx
+ idx = i + (frame_idx < num_frames - 1)
frame = {
- "action": torch.from_numpy(action[i]),
+ "action": action[i],
# Shift reward and success by +1 until the last item of the episode
- "next.reward": reward[i + (frame_idx < num_frames - 1)],
- "next.success": success[i + (frame_idx < num_frames - 1)],
+ "next.reward": reward[idx : idx + 1],
+ "next.success": success[idx : idx + 1],
+ "task": PUSHT_TASK,
}
- frame["observation.state"] = torch.from_numpy(agent_pos[i])
+ frame["observation.state"] = agent_pos[i]
if mode == "keypoints":
- frame["observation.environment_state"] = torch.from_numpy(keypoints[i])
+ frame["observation.environment_state"] = keypoints[i]
else:
- frame["observation.image"] = torch.from_numpy(image[i])
+ frame["observation.image"] = image[i]
dataset.add_frame(frame)
- dataset.save_episode(task=PUSHT_TASK)
-
- dataset.consolidate()
+ dataset.save_episode()
if push_to_hub:
dataset.push_to_hub()
+ hub_api = HfApi()
+ hub_api.create_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
if __name__ == "__main__":
@@ -218,5 +239,5 @@ if __name__ == "__main__":
main(raw_dir, repo_id=repo_id, mode=mode)
# Uncomment if you want to load the local dataset and explore it
- # dataset = LeRobotDataset(repo_id=repo_id, local_files_only=True)
+ # dataset = LeRobotDataset(repo_id=repo_id)
# breakpoint()
diff --git a/lerobot/common/constants.py b/lerobot/common/constants.py
index 73889594..973595cd 100644
--- a/lerobot/common/constants.py
+++ b/lerobot/common/constants.py
@@ -1,6 +1,45 @@
+# 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.
# keys
+import os
+from pathlib import Path
+
+from huggingface_hub.constants import HF_HOME
+
OBS_ENV = "observation.environment_state"
OBS_ROBOT = "observation.state"
OBS_IMAGE = "observation.image"
OBS_IMAGES = "observation.images"
ACTION = "action"
+
+# files & directories
+CHECKPOINTS_DIR = "checkpoints"
+LAST_CHECKPOINT_LINK = "last"
+PRETRAINED_MODEL_DIR = "pretrained_model"
+TRAINING_STATE_DIR = "training_state"
+RNG_STATE = "rng_state.safetensors"
+TRAINING_STEP = "training_step.json"
+OPTIMIZER_STATE = "optimizer_state.safetensors"
+OPTIMIZER_PARAM_GROUPS = "optimizer_param_groups.json"
+SCHEDULER_STATE = "scheduler_state.json"
+
+# cache dir
+default_cache_path = Path(HF_HOME) / "lerobot"
+HF_LEROBOT_HOME = Path(os.getenv("HF_LEROBOT_HOME", default_cache_path)).expanduser()
+
+if "LEROBOT_HOME" in os.environ:
+ raise ValueError(
+ f"You have a 'LEROBOT_HOME' environment variable set to '{os.getenv('LEROBOT_HOME')}'.\n"
+ "'LEROBOT_HOME' is deprecated, please use 'HF_LEROBOT_HOME' instead."
+ )
diff --git a/lerobot/common/datasets/backward_compatibility.py b/lerobot/common/datasets/backward_compatibility.py
new file mode 100644
index 00000000..cf8e31c4
--- /dev/null
+++ b/lerobot/common/datasets/backward_compatibility.py
@@ -0,0 +1,68 @@
+# 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 packaging.version
+
+V2_MESSAGE = """
+The dataset you requested ({repo_id}) is in {version} format.
+
+We introduced a new format since v2.0 which is not backward compatible with v1.x.
+Please, use our conversion script. Modify the following command with your own task description:
+```
+python lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py \\
+ --repo-id {repo_id} \\
+ --single-task "TASK DESCRIPTION." # <---- /!\\ Replace TASK DESCRIPTION /!\\
+```
+
+A few examples to replace TASK DESCRIPTION: "Pick up the blue cube and place it into the bin.", "Insert the
+peg into the socket.", "Slide open the ziploc bag.", "Take the elevator to the 1st floor.", "Open the top
+cabinet, store the pot inside it then close the cabinet.", "Push the T-shaped block onto the T-shaped
+target.", "Grab the spray paint on the shelf and place it in the bin on top of the robot dog.", "Fold the
+sweatshirt.", ...
+
+If you encounter a problem, contact LeRobot maintainers on [Discord](https://discord.com/invite/s3KuuzsPFb)
+or open an [issue on GitHub](https://github.com/huggingface/lerobot/issues/new/choose).
+"""
+
+V21_MESSAGE = """
+The dataset you requested ({repo_id}) is in {version} format.
+While current version of LeRobot is backward-compatible with it, the version of your dataset still uses global
+stats instead of per-episode stats. Update your dataset stats to the new format using this command:
+```
+python lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py --repo-id={repo_id}
+```
+
+If you encounter a problem, contact LeRobot maintainers on [Discord](https://discord.com/invite/s3KuuzsPFb)
+or open an [issue on GitHub](https://github.com/huggingface/lerobot/issues/new/choose).
+"""
+
+FUTURE_MESSAGE = """
+The dataset you requested ({repo_id}) is only available in {version} format.
+As we cannot ensure forward compatibility with it, please update your current version of lerobot.
+"""
+
+
+class CompatibilityError(Exception): ...
+
+
+class BackwardCompatibilityError(CompatibilityError):
+ def __init__(self, repo_id: str, version: packaging.version.Version):
+ message = V2_MESSAGE.format(repo_id=repo_id, version=version)
+ super().__init__(message)
+
+
+class ForwardCompatibilityError(CompatibilityError):
+ def __init__(self, repo_id: str, version: packaging.version.Version):
+ message = FUTURE_MESSAGE.format(repo_id=repo_id, version=version)
+ super().__init__(message)
diff --git a/lerobot/common/datasets/compute_stats.py b/lerobot/common/datasets/compute_stats.py
index c6211699..1149ec83 100644
--- a/lerobot/common/datasets/compute_stats.py
+++ b/lerobot/common/datasets/compute_stats.py
@@ -13,202 +13,164 @@
# 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.
-from copy import deepcopy
-from math import ceil
+import numpy as np
-import einops
-import torch
-import tqdm
+from lerobot.common.datasets.utils import load_image_as_numpy
-def get_stats_einops_patterns(dataset, num_workers=0):
- """These einops patterns will be used to aggregate batches and compute statistics.
+def estimate_num_samples(
+ dataset_len: int, min_num_samples: int = 100, max_num_samples: int = 10_000, power: float = 0.75
+) -> int:
+ """Heuristic to estimate the number of samples based on dataset size.
+ The power controls the sample growth relative to dataset size.
+ Lower the power for less number of samples.
- Note: We assume the images are in channel first format
+ For default arguments, we have:
+ - from 1 to ~500, num_samples=100
+ - at 1000, num_samples=177
+ - at 2000, num_samples=299
+ - at 5000, num_samples=594
+ - at 10000, num_samples=1000
+ - at 20000, num_samples=1681
"""
+ if dataset_len < min_num_samples:
+ min_num_samples = dataset_len
+ return max(min_num_samples, min(int(dataset_len**power), max_num_samples))
- dataloader = torch.utils.data.DataLoader(
- dataset,
- num_workers=num_workers,
- batch_size=2,
- shuffle=False,
- )
- batch = next(iter(dataloader))
- stats_patterns = {}
+def sample_indices(data_len: int) -> list[int]:
+ num_samples = estimate_num_samples(data_len)
+ return np.round(np.linspace(0, data_len - 1, num_samples)).astype(int).tolist()
- for key in dataset.features:
- # sanity check that tensors are not float64
- assert batch[key].dtype != torch.float64
- # if isinstance(feats_type, (VideoFrame, Image)):
- if key in dataset.meta.camera_keys:
- # sanity check that images are channel first
- _, c, h, w = batch[key].shape
- assert c < h and c < w, f"expect channel first images, but instead {batch[key].shape}"
+def auto_downsample_height_width(img: np.ndarray, target_size: int = 150, max_size_threshold: int = 300):
+ _, height, width = img.shape
- # sanity check that images are float32 in range [0,1]
- assert batch[key].dtype == torch.float32, f"expect torch.float32, but instead {batch[key].dtype=}"
- assert batch[key].max() <= 1, f"expect pixels lower than 1, but instead {batch[key].max()=}"
- assert batch[key].min() >= 0, f"expect pixels greater than 1, but instead {batch[key].min()=}"
+ if max(width, height) < max_size_threshold:
+ # no downsampling needed
+ return img
- stats_patterns[key] = "b c h w -> c 1 1"
- elif batch[key].ndim == 2:
- stats_patterns[key] = "b c -> c "
- elif batch[key].ndim == 1:
- stats_patterns[key] = "b -> 1"
+ downsample_factor = int(width / target_size) if width > height else int(height / target_size)
+ return img[:, ::downsample_factor, ::downsample_factor]
+
+
+def sample_images(image_paths: list[str]) -> np.ndarray:
+ sampled_indices = sample_indices(len(image_paths))
+
+ images = None
+ for i, idx in enumerate(sampled_indices):
+ path = image_paths[idx]
+ # we load as uint8 to reduce memory usage
+ img = load_image_as_numpy(path, dtype=np.uint8, channel_first=True)
+ img = auto_downsample_height_width(img)
+
+ if images is None:
+ images = np.empty((len(sampled_indices), *img.shape), dtype=np.uint8)
+
+ images[i] = img
+
+ return images
+
+
+def get_feature_stats(array: np.ndarray, axis: tuple, keepdims: bool) -> dict[str, np.ndarray]:
+ return {
+ "min": np.min(array, axis=axis, keepdims=keepdims),
+ "max": np.max(array, axis=axis, keepdims=keepdims),
+ "mean": np.mean(array, axis=axis, keepdims=keepdims),
+ "std": np.std(array, axis=axis, keepdims=keepdims),
+ "count": np.array([len(array)]),
+ }
+
+
+def compute_episode_stats(episode_data: dict[str, list[str] | np.ndarray], features: dict) -> dict:
+ ep_stats = {}
+ for key, data in episode_data.items():
+ if features[key]["dtype"] == "string":
+ continue # HACK: we should receive np.arrays of strings
+ elif features[key]["dtype"] in ["image", "video"]:
+ ep_ft_array = sample_images(data) # data is a list of image paths
+ axes_to_reduce = (0, 2, 3) # keep channel dim
+ keepdims = True
else:
- raise ValueError(f"{key}, {batch[key].shape}")
+ ep_ft_array = data # data is already a np.ndarray
+ axes_to_reduce = 0 # compute stats over the first axis
+ keepdims = data.ndim == 1 # keep as np.array
- return stats_patterns
+ ep_stats[key] = get_feature_stats(ep_ft_array, axis=axes_to_reduce, keepdims=keepdims)
+
+ # finally, we normalize and remove batch dim for images
+ if features[key]["dtype"] in ["image", "video"]:
+ ep_stats[key] = {
+ k: v if k == "count" else np.squeeze(v / 255.0, axis=0) for k, v in ep_stats[key].items()
+ }
+
+ return ep_stats
-def compute_stats(dataset, batch_size=8, num_workers=8, max_num_samples=None):
- """Compute mean/std and min/max statistics of all data keys in a LeRobotDataset."""
- if max_num_samples is None:
- max_num_samples = len(dataset)
-
- # for more info on why we need to set the same number of workers, see `load_from_videos`
- stats_patterns = get_stats_einops_patterns(dataset, num_workers)
-
- # mean and std will be computed incrementally while max and min will track the running value.
- mean, std, max, min = {}, {}, {}, {}
- for key in stats_patterns:
- mean[key] = torch.tensor(0.0).float()
- std[key] = torch.tensor(0.0).float()
- max[key] = torch.tensor(-float("inf")).float()
- min[key] = torch.tensor(float("inf")).float()
-
- def create_seeded_dataloader(dataset, batch_size, seed):
- generator = torch.Generator()
- generator.manual_seed(seed)
- dataloader = torch.utils.data.DataLoader(
- dataset,
- num_workers=num_workers,
- batch_size=batch_size,
- shuffle=True,
- drop_last=False,
- generator=generator,
- )
- return dataloader
-
- # Note: Due to be refactored soon. The point of storing `first_batch` is to make sure we don't get
- # surprises when rerunning the sampler.
- first_batch = None
- running_item_count = 0 # for online mean computation
- dataloader = create_seeded_dataloader(dataset, batch_size, seed=1337)
- for i, batch in enumerate(
- tqdm.tqdm(dataloader, total=ceil(max_num_samples / batch_size), desc="Compute mean, min, max")
- ):
- this_batch_size = len(batch["index"])
- running_item_count += this_batch_size
- if first_batch is None:
- first_batch = deepcopy(batch)
- for key, pattern in stats_patterns.items():
- batch[key] = batch[key].float()
- # Numerically stable update step for mean computation.
- batch_mean = einops.reduce(batch[key], pattern, "mean")
- # Hint: to update the mean we need x̄ₙ = (Nₙ₋₁x̄ₙ₋₁ + Bₙxₙ) / Nₙ, where the subscript represents
- # the update step, N is the running item count, B is this batch size, x̄ is the running mean,
- # and x is the current batch mean. Some rearrangement is then required to avoid risking
- # numerical overflow. Another hint: Nₙ₋₁ = Nₙ - Bₙ. Rearrangement yields
- # x̄ₙ = x̄ₙ₋₁ + Bₙ * (xₙ - x̄ₙ₋₁) / Nₙ
- mean[key] = mean[key] + this_batch_size * (batch_mean - mean[key]) / running_item_count
- max[key] = torch.maximum(max[key], einops.reduce(batch[key], pattern, "max"))
- min[key] = torch.minimum(min[key], einops.reduce(batch[key], pattern, "min"))
-
- if i == ceil(max_num_samples / batch_size) - 1:
- break
-
- first_batch_ = None
- running_item_count = 0 # for online std computation
- dataloader = create_seeded_dataloader(dataset, batch_size, seed=1337)
- for i, batch in enumerate(
- tqdm.tqdm(dataloader, total=ceil(max_num_samples / batch_size), desc="Compute std")
- ):
- this_batch_size = len(batch["index"])
- running_item_count += this_batch_size
- # Sanity check to make sure the batches are still in the same order as before.
- if first_batch_ is None:
- first_batch_ = deepcopy(batch)
- for key in stats_patterns:
- assert torch.equal(first_batch_[key], first_batch[key])
- for key, pattern in stats_patterns.items():
- batch[key] = batch[key].float()
- # Numerically stable update step for mean computation (where the mean is over squared
- # residuals).See notes in the mean computation loop above.
- batch_std = einops.reduce((batch[key] - mean[key]) ** 2, pattern, "mean")
- std[key] = std[key] + this_batch_size * (batch_std - std[key]) / running_item_count
-
- if i == ceil(max_num_samples / batch_size) - 1:
- break
-
- for key in stats_patterns:
- std[key] = torch.sqrt(std[key])
-
- stats = {}
- for key in stats_patterns:
- stats[key] = {
- "mean": mean[key],
- "std": std[key],
- "max": max[key],
- "min": min[key],
- }
- return stats
+def _assert_type_and_shape(stats_list: list[dict[str, dict]]):
+ for i in range(len(stats_list)):
+ for fkey in stats_list[i]:
+ for k, v in stats_list[i][fkey].items():
+ if not isinstance(v, np.ndarray):
+ raise ValueError(
+ f"Stats must be composed of numpy array, but key '{k}' of feature '{fkey}' is of type '{type(v)}' instead."
+ )
+ if v.ndim == 0:
+ raise ValueError("Number of dimensions must be at least 1, and is 0 instead.")
+ if k == "count" and v.shape != (1,):
+ raise ValueError(f"Shape of 'count' must be (1), but is {v.shape} instead.")
+ if "image" in fkey and k != "count" and v.shape != (3, 1, 1):
+ raise ValueError(f"Shape of '{k}' must be (3,1,1), but is {v.shape} instead.")
-def aggregate_stats(ls_datasets) -> dict[str, torch.Tensor]:
- """Aggregate stats of multiple LeRobot datasets into one set of stats without recomputing from scratch.
+def aggregate_feature_stats(stats_ft_list: list[dict[str, dict]]) -> dict[str, dict[str, np.ndarray]]:
+ """Aggregates stats for a single feature."""
+ means = np.stack([s["mean"] for s in stats_ft_list])
+ variances = np.stack([s["std"] ** 2 for s in stats_ft_list])
+ counts = np.stack([s["count"] for s in stats_ft_list])
+ total_count = counts.sum(axis=0)
- The final stats will have the union of all data keys from each of the datasets.
+ # Prepare weighted mean by matching number of dimensions
+ while counts.ndim < means.ndim:
+ counts = np.expand_dims(counts, axis=-1)
- The final stats will have the union of all data keys from each of the datasets. For instance:
- - new_max = max(max_dataset_0, max_dataset_1, ...)
+ # Compute the weighted mean
+ weighted_means = means * counts
+ total_mean = weighted_means.sum(axis=0) / total_count
+
+ # Compute the variance using the parallel algorithm
+ delta_means = means - total_mean
+ weighted_variances = (variances + delta_means**2) * counts
+ total_variance = weighted_variances.sum(axis=0) / total_count
+
+ return {
+ "min": np.min(np.stack([s["min"] for s in stats_ft_list]), axis=0),
+ "max": np.max(np.stack([s["max"] for s in stats_ft_list]), axis=0),
+ "mean": total_mean,
+ "std": np.sqrt(total_variance),
+ "count": total_count,
+ }
+
+
+def aggregate_stats(stats_list: list[dict[str, dict]]) -> dict[str, dict[str, np.ndarray]]:
+ """Aggregate stats from multiple compute_stats outputs into a single set of stats.
+
+ The final stats will have the union of all data keys from each of the stats dicts.
+
+ For instance:
- new_min = min(min_dataset_0, min_dataset_1, ...)
- - new_mean = (mean of all data)
+ - new_max = max(max_dataset_0, max_dataset_1, ...)
+ - new_mean = (mean of all data, weighted by counts)
- new_std = (std of all data)
"""
- data_keys = set()
- for dataset in ls_datasets:
- data_keys.update(dataset.meta.stats.keys())
- stats = {k: {} for k in data_keys}
- for data_key in data_keys:
- for stat_key in ["min", "max"]:
- # compute `max(dataset_0["max"], dataset_1["max"], ...)`
- stats[data_key][stat_key] = einops.reduce(
- torch.stack(
- [ds.meta.stats[data_key][stat_key] for ds in ls_datasets if data_key in ds.meta.stats],
- dim=0,
- ),
- "n ... -> ...",
- stat_key,
- )
- total_samples = sum(d.num_frames for d in ls_datasets if data_key in d.meta.stats)
- # Compute the "sum" statistic by multiplying each mean by the number of samples in the respective
- # dataset, then divide by total_samples to get the overall "mean".
- # NOTE: the brackets around (d.num_frames / total_samples) are needed tor minimize the risk of
- # numerical overflow!
- stats[data_key]["mean"] = sum(
- d.meta.stats[data_key]["mean"] * (d.num_frames / total_samples)
- for d in ls_datasets
- if data_key in d.meta.stats
- )
- # The derivation for standard deviation is a little more involved but is much in the same spirit as
- # the computation of the mean.
- # Given two sets of data where the statistics are known:
- # σ_combined = sqrt[ (n1 * (σ1^2 + d1^2) + n2 * (σ2^2 + d2^2)) / (n1 + n2) ]
- # where d1 = μ1 - μ_combined, d2 = μ2 - μ_combined
- # NOTE: the brackets around (d.num_frames / total_samples) are needed tor minimize the risk of
- # numerical overflow!
- stats[data_key]["std"] = torch.sqrt(
- sum(
- (
- d.meta.stats[data_key]["std"] ** 2
- + (d.meta.stats[data_key]["mean"] - stats[data_key]["mean"]) ** 2
- )
- * (d.num_frames / total_samples)
- for d in ls_datasets
- if data_key in d.meta.stats
- )
- )
- return stats
+
+ _assert_type_and_shape(stats_list)
+
+ data_keys = {key for stats in stats_list for key in stats}
+ aggregated_stats = {key: {} for key in data_keys}
+
+ for key in data_keys:
+ stats_with_key = [stats[key] for stats in stats_list if key in stats]
+ aggregated_stats[key] = aggregate_feature_stats(stats_with_key)
+
+ return aggregated_stats
diff --git a/lerobot/common/datasets/factory.py b/lerobot/common/datasets/factory.py
index 58ff400e..38c01b42 100644
--- a/lerobot/common/datasets/factory.py
+++ b/lerobot/common/datasets/factory.py
@@ -83,15 +83,18 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
)
if isinstance(cfg.dataset.repo_id, str):
- ds_meta = LeRobotDatasetMetadata(cfg.dataset.repo_id, local_files_only=cfg.dataset.local_files_only)
+ ds_meta = LeRobotDatasetMetadata(
+ cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
+ )
delta_timestamps = resolve_delta_timestamps(cfg.policy, ds_meta)
dataset = LeRobotDataset(
cfg.dataset.repo_id,
+ root=cfg.dataset.root,
episodes=cfg.dataset.episodes,
delta_timestamps=delta_timestamps,
image_transforms=image_transforms,
+ revision=cfg.dataset.revision,
video_backend=cfg.dataset.video_backend,
- local_files_only=cfg.dataset.local_files_only,
)
else:
raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.")
@@ -104,7 +107,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
)
logging.info(
"Multiple datasets were provided. Applied the following index mapping to the provided datasets: "
- f"{pformat(dataset.repo_id_to_index , indent=2)}"
+ f"{pformat(dataset.repo_id_to_index, indent=2)}"
)
if cfg.dataset.use_imagenet_stats:
diff --git a/lerobot/common/datasets/image_writer.py b/lerobot/common/datasets/image_writer.py
index 85dd6830..6fc0ee2f 100644
--- a/lerobot/common/datasets/image_writer.py
+++ b/lerobot/common/datasets/image_writer.py
@@ -38,22 +38,40 @@ def safe_stop_image_writer(func):
return wrapper
-def image_array_to_image(image_array: np.ndarray) -> PIL.Image.Image:
+def image_array_to_pil_image(image_array: np.ndarray, range_check: bool = True) -> PIL.Image.Image:
# TODO(aliberts): handle 1 channel and 4 for depth images
- if image_array.ndim == 3 and image_array.shape[0] in [1, 3]:
+ if image_array.ndim != 3:
+ raise ValueError(f"The array has {image_array.ndim} dimensions, but 3 is expected for an image.")
+
+ if image_array.shape[0] == 3:
# Transpose from pytorch convention (C, H, W) to (H, W, C)
image_array = image_array.transpose(1, 2, 0)
+
+ elif image_array.shape[-1] != 3:
+ raise NotImplementedError(
+ f"The image has {image_array.shape[-1]} channels, but 3 is required for now."
+ )
+
if image_array.dtype != np.uint8:
- # Assume the image is in [0, 1] range for floating-point data
- image_array = np.clip(image_array, 0, 1)
+ if range_check:
+ max_ = image_array.max().item()
+ min_ = image_array.min().item()
+ if max_ > 1.0 or min_ < 0.0:
+ raise ValueError(
+ "The image data type is float, which requires values in the range [0.0, 1.0]. "
+ f"However, the provided range is [{min_}, {max_}]. Please adjust the range or "
+ "provide a uint8 image with values in the range [0, 255]."
+ )
+
image_array = (image_array * 255).astype(np.uint8)
+
return PIL.Image.fromarray(image_array)
def write_image(image: np.ndarray | PIL.Image.Image, fpath: Path):
try:
if isinstance(image, np.ndarray):
- img = image_array_to_image(image)
+ img = image_array_to_pil_image(image)
elif isinstance(image, PIL.Image.Image):
img = image
else:
diff --git a/lerobot/common/datasets/lerobot_dataset.py b/lerobot/common/datasets/lerobot_dataset.py
index 9483bf0a..6ef955dd 100644
--- a/lerobot/common/datasets/lerobot_dataset.py
+++ b/lerobot/common/datasets/lerobot_dataset.py
@@ -13,62 +13,68 @@
# 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 contextlib
import logging
-import os
import shutil
-from functools import cached_property
from pathlib import Path
from typing import Callable
import datasets
import numpy as np
+import packaging.version
import PIL.Image
import torch
import torch.utils
-from datasets import load_dataset
-from huggingface_hub import create_repo, snapshot_download, upload_folder
+from datasets import concatenate_datasets, load_dataset
+from huggingface_hub import HfApi, snapshot_download
+from huggingface_hub.constants import REPOCARD_NAME
+from huggingface_hub.errors import RevisionNotFoundError
-from lerobot.common.datasets.compute_stats import aggregate_stats, compute_stats
+from lerobot.common.constants import HF_LEROBOT_HOME
+from lerobot.common.datasets.compute_stats import aggregate_stats, compute_episode_stats
from lerobot.common.datasets.image_writer import AsyncImageWriter, write_image
from lerobot.common.datasets.utils import (
DEFAULT_FEATURES,
DEFAULT_IMAGE_PATH,
- EPISODES_PATH,
INFO_PATH,
- STATS_PATH,
TASKS_PATH,
append_jsonlines,
+ backward_compatible_episodes_stats,
check_delta_timestamps,
check_timestamps_sync,
check_version_compatibility,
- create_branch,
create_empty_dataset_info,
create_lerobot_dataset_card,
+ embed_images,
get_delta_indices,
get_episode_data_index,
get_features_from_robot,
get_hf_features_from_features,
- get_hub_safe_version,
+ get_safe_version,
hf_transform_to_torch,
+ is_valid_version,
load_episodes,
+ load_episodes_stats,
load_info,
load_stats,
load_tasks,
- serialize_dict,
+ validate_episode_buffer,
+ validate_frame,
+ write_episode,
+ write_episode_stats,
+ write_info,
write_json,
- write_parquet,
)
from lerobot.common.datasets.video_utils import (
VideoFrame,
- decode_video_frames_torchvision,
+ decode_video_frames,
encode_video_frames,
+ get_safe_default_codec,
get_video_info,
)
from lerobot.common.robot_devices.robots.utils import Robot
-# For maintainers, see lerobot/common/datasets/push_dataset_to_hub/CODEBASE_VERSION.md
-CODEBASE_VERSION = "v2.0"
-LEROBOT_HOME = Path(os.getenv("LEROBOT_HOME", "~/.cache/huggingface/lerobot")).expanduser()
+CODEBASE_VERSION = "v2.1"
class LeRobotDatasetMetadata:
@@ -76,19 +82,36 @@ class LeRobotDatasetMetadata:
self,
repo_id: str,
root: str | Path | None = None,
- local_files_only: bool = False,
+ revision: str | None = None,
+ force_cache_sync: bool = False,
):
self.repo_id = repo_id
- self.root = Path(root) if root is not None else LEROBOT_HOME / repo_id
- self.local_files_only = local_files_only
+ self.revision = revision if revision else CODEBASE_VERSION
+ self.root = Path(root) if root is not None else HF_LEROBOT_HOME / repo_id
- # Load metadata
- (self.root / "meta").mkdir(exist_ok=True, parents=True)
- self.pull_from_repo(allow_patterns="meta/")
+ try:
+ if force_cache_sync:
+ raise FileNotFoundError
+ self.load_metadata()
+ except (FileNotFoundError, NotADirectoryError):
+ if is_valid_version(self.revision):
+ self.revision = get_safe_version(self.repo_id, self.revision)
+
+ (self.root / "meta").mkdir(exist_ok=True, parents=True)
+ self.pull_from_repo(allow_patterns="meta/")
+ self.load_metadata()
+
+ def load_metadata(self):
self.info = load_info(self.root)
- self.stats = load_stats(self.root)
- self.tasks = load_tasks(self.root)
+ check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
+ self.tasks, self.task_to_task_index = load_tasks(self.root)
self.episodes = load_episodes(self.root)
+ if self._version < packaging.version.parse("v2.1"):
+ self.stats = load_stats(self.root)
+ self.episodes_stats = backward_compatible_episodes_stats(self.stats, self.episodes)
+ else:
+ self.episodes_stats = load_episodes_stats(self.root)
+ self.stats = aggregate_stats(list(self.episodes_stats.values()))
def pull_from_repo(
self,
@@ -98,21 +121,16 @@ class LeRobotDatasetMetadata:
snapshot_download(
self.repo_id,
repo_type="dataset",
- revision=self._hub_version,
+ revision=self.revision,
local_dir=self.root,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
- local_files_only=self.local_files_only,
)
- @cached_property
- def _hub_version(self) -> str | None:
- return None if self.local_files_only else get_hub_safe_version(self.repo_id, CODEBASE_VERSION)
-
@property
- def _version(self) -> str:
+ def _version(self) -> packaging.version.Version:
"""Codebase version used to create this dataset."""
- return self.info["codebase_version"]
+ return packaging.version.parse(self.info["codebase_version"])
def get_data_file_path(self, ep_index: int) -> Path:
ep_chunk = self.get_episode_chunk(ep_index)
@@ -202,54 +220,65 @@ class LeRobotDatasetMetadata:
"""Max number of episodes per chunk."""
return self.info["chunks_size"]
- @property
- def task_to_task_index(self) -> dict:
- return {task: task_idx for task_idx, task in self.tasks.items()}
-
- def get_task_index(self, task: str) -> int:
+ def get_task_index(self, task: str) -> int | None:
"""
Given a task in natural language, returns its task_index if the task already exists in the dataset,
- otherwise creates a new task_index.
+ otherwise return None.
"""
- task_index = self.task_to_task_index.get(task, None)
- return task_index if task_index is not None else self.total_tasks
+ return self.task_to_task_index.get(task, None)
- def save_episode(self, episode_index: int, episode_length: int, task: str, task_index: int) -> None:
+ def add_task(self, task: str):
+ """
+ Given a task in natural language, add it to the dictionary of tasks.
+ """
+ if task in self.task_to_task_index:
+ raise ValueError(f"The task '{task}' already exists and can't be added twice.")
+
+ task_index = self.info["total_tasks"]
+ self.task_to_task_index[task] = task_index
+ self.tasks[task_index] = task
+ self.info["total_tasks"] += 1
+
+ task_dict = {
+ "task_index": task_index,
+ "task": task,
+ }
+ append_jsonlines(task_dict, self.root / TASKS_PATH)
+
+ def save_episode(
+ self,
+ episode_index: int,
+ episode_length: int,
+ episode_tasks: list[str],
+ episode_stats: dict[str, dict],
+ ) -> None:
self.info["total_episodes"] += 1
self.info["total_frames"] += episode_length
- if task_index not in self.tasks:
- self.info["total_tasks"] += 1
- self.tasks[task_index] = task
- task_dict = {
- "task_index": task_index,
- "task": task,
- }
- append_jsonlines(task_dict, self.root / TASKS_PATH)
-
chunk = self.get_episode_chunk(episode_index)
if chunk >= self.total_chunks:
self.info["total_chunks"] += 1
self.info["splits"] = {"train": f"0:{self.info['total_episodes']}"}
self.info["total_videos"] += len(self.video_keys)
- write_json(self.info, self.root / INFO_PATH)
+ if len(self.video_keys) > 0:
+ self.update_video_info()
+
+ write_info(self.info, self.root)
episode_dict = {
"episode_index": episode_index,
- "tasks": [task],
+ "tasks": episode_tasks,
"length": episode_length,
}
- self.episodes.append(episode_dict)
- append_jsonlines(episode_dict, self.root / EPISODES_PATH)
+ self.episodes[episode_index] = episode_dict
+ write_episode(episode_dict, self.root)
- # TODO(aliberts): refactor stats in save_episodes
- # image_sampling = int(self.fps / 2) # sample 2 img/s for the stats
- # ep_stats = compute_episode_stats(episode_buffer, self.features, episode_length, image_sampling=image_sampling)
- # ep_stats = serialize_dict(ep_stats)
- # append_jsonlines(ep_stats, self.root / STATS_PATH)
+ self.episodes_stats[episode_index] = episode_stats
+ self.stats = aggregate_stats([self.stats, episode_stats]) if self.stats else episode_stats
+ write_episode_stats(episode_index, episode_stats, self.root)
- def write_video_info(self) -> None:
+ def update_video_info(self) -> None:
"""
Warning: this function writes info from first episode videos, implicitly assuming that all videos have
been encoded the same way. Also, this means it assumes the first episode exists.
@@ -259,8 +288,6 @@ class LeRobotDatasetMetadata:
video_path = self.root / self.get_video_file_path(ep_index=0, vid_key=key)
self.info["features"][key]["info"] = get_video_info(video_path)
- write_json(self.info, self.root / INFO_PATH)
-
def __repr__(self):
feature_keys = list(self.features)
return (
@@ -286,7 +313,7 @@ class LeRobotDatasetMetadata:
"""Creates metadata for a LeRobotDataset."""
obj = cls.__new__(cls)
obj.repo_id = repo_id
- obj.root = Path(root) if root is not None else LEROBOT_HOME / repo_id
+ obj.root = Path(root) if root is not None else HF_LEROBOT_HOME / repo_id
obj.root.mkdir(parents=True, exist_ok=False)
@@ -304,6 +331,7 @@ class LeRobotDatasetMetadata:
)
else:
# TODO(aliberts, rcadene): implement sanity check for features
+ features = {**features, **DEFAULT_FEATURES}
# check if none of the features contains a "/" in their names,
# as this would break the dict flattening in the stats computation, which uses '/' as separator
@@ -313,12 +341,13 @@ class LeRobotDatasetMetadata:
features = {**features, **DEFAULT_FEATURES}
- obj.tasks, obj.stats, obj.episodes = {}, {}, []
+ obj.tasks, obj.task_to_task_index = {}, {}
+ obj.episodes_stats, obj.stats, obj.episodes = {}, {}, {}
obj.info = create_empty_dataset_info(CODEBASE_VERSION, fps, robot_type, features, use_videos)
if len(obj.video_keys) > 0 and not use_videos:
raise ValueError()
write_json(obj.info, obj.root / INFO_PATH)
- obj.local_files_only = True
+ obj.revision = None
return obj
@@ -331,8 +360,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
image_transforms: Callable | None = None,
delta_timestamps: dict[list[float]] | None = None,
tolerance_s: float = 1e-4,
+ revision: str | None = None,
+ force_cache_sync: bool = False,
download_videos: bool = True,
- local_files_only: bool = False,
video_backend: str | None = None,
):
"""
@@ -342,7 +372,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
- On your local disk in the 'root' folder. This is typically the case when you recorded your
dataset locally and you may or may not have pushed it to the hub yet. Instantiating this class
with 'root' will load your dataset directly from disk. This can happen while you're offline (no
- internet connection), in that case, use local_files_only=True.
+ internet connection).
- On the Hugging Face Hub at the address https://huggingface.co/datasets/{repo_id} and not on
your local disk in the 'root' folder. Instantiating this class with this 'repo_id' will download
@@ -362,7 +392,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
- info contains various information about the dataset like shapes, keys, fps etc.
- stats stores the dataset statistics of the different modalities for normalization
- tasks contains the prompts for each task of the dataset, which can be used for
- task-conditionned training.
+ task-conditioned training.
- hf_dataset (from datasets.Dataset), which will read any values from parquet files.
- videos (optional) from which frames are loaded to be synchronous with data from parquet files.
@@ -424,24 +454,28 @@ class LeRobotDataset(torch.utils.data.Dataset):
timestamps is separated to the next by 1/fps +/- tolerance_s. This also applies to frames
decoded from video files. It is also used to check that `delta_timestamps` (when provided) are
multiples of 1/fps. Defaults to 1e-4.
+ revision (str, optional): An optional Git revision id which can be a branch name, a tag, or a
+ commit hash. Defaults to current codebase version tag.
+ sync_cache_first (bool, optional): Flag to sync and refresh local files first. If True and files
+ are already present in the local cache, this will be faster. However, files loaded might not
+ be in sync with the version on the hub, especially if you specified 'revision'. Defaults to
+ False.
download_videos (bool, optional): Flag to download the videos. Note that when set to True but the
video files are already present on local disk, they won't be downloaded again. Defaults to
True.
- local_files_only (bool, optional): Flag to use local files only. If True, no requests to the hub
- will be made. Defaults to False.
- video_backend (str | None, optional): Video backend to use for decoding videos. There is currently
- a single option which is the pyav decoder used by Torchvision. Defaults to pyav.
+ video_backend (str | None, optional): Video backend to use for decoding videos. Defaults to torchcodec when available int the platform; otherwise, defaults to 'pyav'.
+ You can also use the 'pyav' decoder used by Torchvision, which used to be the default option, or 'video_reader' which is another decoder of Torchvision.
"""
super().__init__()
self.repo_id = repo_id
- self.root = Path(root) if root else LEROBOT_HOME / repo_id
+ self.root = Path(root) if root else HF_LEROBOT_HOME / repo_id
self.image_transforms = image_transforms
self.delta_timestamps = delta_timestamps
self.episodes = episodes
self.tolerance_s = tolerance_s
- self.video_backend = video_backend if video_backend else "pyav"
+ self.revision = revision if revision else CODEBASE_VERSION
+ self.video_backend = video_backend if video_backend else get_safe_default_codec()
self.delta_indices = None
- self.local_files_only = local_files_only
# Unused attributes
self.image_writer = None
@@ -450,64 +484,92 @@ class LeRobotDataset(torch.utils.data.Dataset):
self.root.mkdir(exist_ok=True, parents=True)
# Load metadata
- self.meta = LeRobotDatasetMetadata(self.repo_id, self.root, self.local_files_only)
-
- # Check version
- check_version_compatibility(self.repo_id, self.meta._version, CODEBASE_VERSION)
+ self.meta = LeRobotDatasetMetadata(
+ self.repo_id, self.root, self.revision, force_cache_sync=force_cache_sync
+ )
+ if self.episodes is not None and self.meta._version >= packaging.version.parse("v2.1"):
+ episodes_stats = [self.meta.episodes_stats[ep_idx] for ep_idx in self.episodes]
+ self.stats = aggregate_stats(episodes_stats)
# Load actual data
- self.download_episodes(download_videos)
- self.hf_dataset = self.load_hf_dataset()
+ try:
+ if force_cache_sync:
+ raise FileNotFoundError
+ assert all((self.root / fpath).is_file() for fpath in self.get_episodes_file_paths())
+ self.hf_dataset = self.load_hf_dataset()
+ except (AssertionError, FileNotFoundError, NotADirectoryError):
+ self.revision = get_safe_version(self.repo_id, self.revision)
+ self.download_episodes(download_videos)
+ self.hf_dataset = self.load_hf_dataset()
+
self.episode_data_index = get_episode_data_index(self.meta.episodes, self.episodes)
# Check timestamps
- check_timestamps_sync(self.hf_dataset, self.episode_data_index, self.fps, self.tolerance_s)
+ timestamps = torch.stack(self.hf_dataset["timestamp"]).numpy()
+ episode_indices = torch.stack(self.hf_dataset["episode_index"]).numpy()
+ ep_data_index_np = {k: t.numpy() for k, t in self.episode_data_index.items()}
+ check_timestamps_sync(timestamps, episode_indices, ep_data_index_np, self.fps, self.tolerance_s)
# Setup delta_indices
if self.delta_timestamps is not None:
check_delta_timestamps(self.delta_timestamps, self.fps, self.tolerance_s)
self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps)
- # Available stats implies all videos have been encoded and dataset is iterable
- self.consolidated = self.meta.stats is not None
-
def push_to_hub(
self,
+ branch: str | None = None,
tags: list | None = None,
license: str | None = "apache-2.0",
+ tag_version: bool = True,
push_videos: bool = True,
private: bool = False,
+ allow_patterns: list[str] | str | None = None,
+ upload_large_folder: bool = False,
**card_kwargs,
) -> None:
- if not self.consolidated:
- logging.warning(
- "You are trying to upload to the hub a LeRobotDataset that has not been consolidated yet. "
- "Consolidating first."
- )
- self.consolidate()
-
ignore_patterns = ["images/"]
if not push_videos:
ignore_patterns.append("videos/")
- create_repo(
+ hub_api = HfApi()
+ hub_api.create_repo(
repo_id=self.repo_id,
private=private,
repo_type="dataset",
exist_ok=True,
)
+ if branch:
+ hub_api.create_branch(
+ repo_id=self.repo_id,
+ branch=branch,
+ revision=self.revision,
+ repo_type="dataset",
+ exist_ok=True,
+ )
- upload_folder(
- repo_id=self.repo_id,
- folder_path=self.root,
- repo_type="dataset",
- ignore_patterns=ignore_patterns,
- )
- card = create_lerobot_dataset_card(
- tags=tags, dataset_info=self.meta.info, license=license, **card_kwargs
- )
- card.push_to_hub(repo_id=self.repo_id, repo_type="dataset")
- create_branch(repo_id=self.repo_id, branch=CODEBASE_VERSION, repo_type="dataset")
+ upload_kwargs = {
+ "repo_id": self.repo_id,
+ "folder_path": self.root,
+ "repo_type": "dataset",
+ "revision": branch,
+ "allow_patterns": allow_patterns,
+ "ignore_patterns": ignore_patterns,
+ }
+ if upload_large_folder:
+ hub_api.upload_large_folder(**upload_kwargs)
+ else:
+ hub_api.upload_folder(**upload_kwargs)
+
+ if not hub_api.file_exists(self.repo_id, REPOCARD_NAME, repo_type="dataset", revision=branch):
+ card = create_lerobot_dataset_card(
+ tags=tags, dataset_info=self.meta.info, license=license, **card_kwargs
+ )
+ card.push_to_hub(repo_id=self.repo_id, repo_type="dataset", revision=branch)
+
+ if tag_version:
+ with contextlib.suppress(RevisionNotFoundError):
+ hub_api.delete_tag(self.repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
+ hub_api.create_tag(self.repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
def pull_from_repo(
self,
@@ -517,11 +579,10 @@ class LeRobotDataset(torch.utils.data.Dataset):
snapshot_download(
self.repo_id,
repo_type="dataset",
- revision=self.meta._hub_version,
+ revision=self.revision,
local_dir=self.root,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns,
- local_files_only=self.local_files_only,
)
def download_episodes(self, download_videos: bool = True) -> None:
@@ -535,17 +596,23 @@ class LeRobotDataset(torch.utils.data.Dataset):
files = None
ignore_patterns = None if download_videos else "videos/"
if self.episodes is not None:
- files = [str(self.meta.get_data_file_path(ep_idx)) for ep_idx in self.episodes]
- if len(self.meta.video_keys) > 0 and download_videos:
- video_files = [
- str(self.meta.get_video_file_path(ep_idx, vid_key))
- for vid_key in self.meta.video_keys
- for ep_idx in self.episodes
- ]
- files += video_files
+ files = self.get_episodes_file_paths()
self.pull_from_repo(allow_patterns=files, ignore_patterns=ignore_patterns)
+ def get_episodes_file_paths(self) -> list[Path]:
+ episodes = self.episodes if self.episodes is not None else list(range(self.meta.total_episodes))
+ fpaths = [str(self.meta.get_data_file_path(ep_idx)) for ep_idx in episodes]
+ if len(self.meta.video_keys) > 0:
+ video_files = [
+ str(self.meta.get_video_file_path(ep_idx, vid_key))
+ for vid_key in self.meta.video_keys
+ for ep_idx in episodes
+ ]
+ fpaths += video_files
+
+ return fpaths
+
def load_hf_dataset(self) -> datasets.Dataset:
"""hf_dataset contains all the observations, states, actions, rewards, etc."""
if self.episodes is None:
@@ -557,7 +624,15 @@ class LeRobotDataset(torch.utils.data.Dataset):
# TODO(aliberts): hf_dataset.set_format("torch")
hf_dataset.set_transform(hf_transform_to_torch)
+ return hf_dataset
+ def create_hf_dataset(self) -> datasets.Dataset:
+ features = get_hf_features_from_features(self.features)
+ ft_dict = {col: [] for col in features}
+ hf_dataset = datasets.Dataset.from_dict(ft_dict, features=features, split="train")
+
+ # TODO(aliberts): hf_dataset.set_format("torch")
+ hf_dataset.set_transform(hf_transform_to_torch)
return hf_dataset
@property
@@ -624,7 +699,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
if key not in self.meta.video_keys
}
- def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict:
+ def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict[str, torch.Tensor]:
"""Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
in the main process (e.g. by using a second Dataloader with num_workers=0). It will result in a
Segmentation Fault. This probably happens because a memory reference to the video loader is created in
@@ -633,9 +708,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
item = {}
for vid_key, query_ts in query_timestamps.items():
video_path = self.root / self.meta.get_video_file_path(ep_idx, vid_key)
- frames = decode_video_frames_torchvision(
- video_path, query_ts, self.tolerance_s, self.video_backend
- )
+ frames = decode_video_frames(video_path, query_ts, self.tolerance_s, self.video_backend)
item[vid_key] = frames.squeeze(0)
return item
@@ -654,8 +727,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
query_indices = None
if self.delta_indices is not None:
- current_ep_idx = self.episodes.index(ep_idx) if self.episodes is not None else ep_idx
- query_indices, padding = self._get_query_indices(idx, current_ep_idx)
+ query_indices, padding = self._get_query_indices(idx, ep_idx)
query_result = self._query_hf_dataset(query_indices)
item = {**item, **padding}
for key, val in query_result.items():
@@ -691,10 +763,13 @@ class LeRobotDataset(torch.utils.data.Dataset):
def create_episode_buffer(self, episode_index: int | None = None) -> dict:
current_ep_idx = self.meta.total_episodes if episode_index is None else episode_index
- return {
- "size": 0,
- **{key: current_ep_idx if key == "episode_index" else [] for key in self.features},
- }
+ ep_buffer = {}
+ # size and task are special cases that are not in self.features
+ ep_buffer["size"] = 0
+ ep_buffer["task"] = []
+ for key in self.features:
+ ep_buffer[key] = current_ep_idx if key == "episode_index" else []
+ return ep_buffer
def _get_image_file_path(self, episode_index: int, image_key: str, frame_index: int) -> Path:
fpath = DEFAULT_IMAGE_PATH.format(
@@ -716,25 +791,35 @@ class LeRobotDataset(torch.utils.data.Dataset):
temporary directory — nothing is written to disk. To save those frames, the 'save_episode()' method
then needs to be called.
"""
- # TODO(aliberts, rcadene): Add sanity check for the input, check it's numpy or torch,
- # check the dtype and shape matches, etc.
+ # Convert torch to numpy if needed
+ for name in frame:
+ if isinstance(frame[name], torch.Tensor):
+ frame[name] = frame[name].numpy()
+
+ validate_frame(frame, self.features)
if self.episode_buffer is None:
self.episode_buffer = self.create_episode_buffer()
+ # Automatically add frame_index and timestamp to episode buffer
frame_index = self.episode_buffer["size"]
timestamp = frame.pop("timestamp") if "timestamp" in frame else frame_index / self.fps
self.episode_buffer["frame_index"].append(frame_index)
self.episode_buffer["timestamp"].append(timestamp)
+ # Add frame features to episode_buffer
for key in frame:
- if key not in self.features:
- raise ValueError(key)
+ if key == "task":
+ # Note: we associate the task in natural language to its task index during `save_episode`
+ self.episode_buffer["task"].append(frame["task"])
+ continue
- if self.features[key]["dtype"] not in ["image", "video"]:
- item = frame[key].numpy() if isinstance(frame[key], torch.Tensor) else frame[key]
- self.episode_buffer[key].append(item)
- elif self.features[key]["dtype"] in ["image", "video"]:
+ if key not in self.features:
+ raise ValueError(
+ f"An element of the frame is not in the features. '{key}' not in '{self.features.keys()}'."
+ )
+
+ if self.features[key]["dtype"] in ["image", "video"]:
img_path = self._get_image_file_path(
episode_index=self.episode_buffer["episode_index"], image_key=key, frame_index=frame_index
)
@@ -742,80 +827,95 @@ class LeRobotDataset(torch.utils.data.Dataset):
img_path.parent.mkdir(parents=True, exist_ok=True)
self._save_image(frame[key], img_path)
self.episode_buffer[key].append(str(img_path))
+ else:
+ self.episode_buffer[key].append(frame[key])
self.episode_buffer["size"] += 1
- def save_episode(self, task: str, encode_videos: bool = True, episode_data: dict | None = None) -> None:
+ def save_episode(self, episode_data: dict | None = None) -> None:
"""
- This will save to disk the current episode in self.episode_buffer. Note that since it affects files on
- disk, it sets self.consolidated to False to ensure proper consolidation later on before uploading to
- the hub.
+ This will save to disk the current episode in self.episode_buffer.
- Use 'encode_videos' if you want to encode videos during the saving of this episode. Otherwise,
- you can do it later with dataset.consolidate(). This is to give more flexibility on when to spend
- time for video encoding.
+ Args:
+ episode_data (dict | None, optional): Dict containing the episode data to save. If None, this will
+ save the current episode in self.episode_buffer, which is filled with 'add_frame'. Defaults to
+ None.
"""
if not episode_data:
episode_buffer = self.episode_buffer
+ validate_episode_buffer(episode_buffer, self.meta.total_episodes, self.features)
+
+ # size and task are special cases that won't be added to hf_dataset
episode_length = episode_buffer.pop("size")
+ tasks = episode_buffer.pop("task")
+ episode_tasks = list(set(tasks))
episode_index = episode_buffer["episode_index"]
- if episode_index != self.meta.total_episodes:
- # TODO(aliberts): Add option to use existing episode_index
- raise NotImplementedError(
- "You might have manually provided the episode_buffer with an episode_index that doesn't "
- "match the total number of episodes in the dataset. This is not supported for now."
- )
- if episode_length == 0:
- raise ValueError(
- "You must add one or several frames with `add_frame` before calling `add_episode`."
- )
+ episode_buffer["index"] = np.arange(self.meta.total_frames, self.meta.total_frames + episode_length)
+ episode_buffer["episode_index"] = np.full((episode_length,), episode_index)
- task_index = self.meta.get_task_index(task)
+ # Add new tasks to the tasks dictionary
+ for task in episode_tasks:
+ task_index = self.meta.get_task_index(task)
+ if task_index is None:
+ self.meta.add_task(task)
- if not set(episode_buffer.keys()) == set(self.features):
- raise ValueError()
+ # Given tasks in natural language, find their corresponding task indices
+ episode_buffer["task_index"] = np.array([self.meta.get_task_index(task) for task in tasks])
for key, ft in self.features.items():
- if key == "index":
- episode_buffer[key] = np.arange(
- self.meta.total_frames, self.meta.total_frames + episode_length
- )
- elif key == "episode_index":
- episode_buffer[key] = np.full((episode_length,), episode_index)
- elif key == "task_index":
- episode_buffer[key] = np.full((episode_length,), task_index)
- elif ft["dtype"] in ["image", "video"]:
+ # index, episode_index, task_index are already processed above, and image and video
+ # are processed separately by storing image path and frame info as meta data
+ if key in ["index", "episode_index", "task_index"] or ft["dtype"] in ["image", "video"]:
continue
- elif len(ft["shape"]) == 1 and ft["shape"][0] == 1:
- episode_buffer[key] = np.array(episode_buffer[key], dtype=ft["dtype"])
- elif len(ft["shape"]) == 1 and ft["shape"][0] > 1:
- episode_buffer[key] = np.stack(episode_buffer[key])
- else:
- raise ValueError(key)
+ episode_buffer[key] = np.stack(episode_buffer[key])
self._wait_image_writer()
self._save_episode_table(episode_buffer, episode_index)
+ ep_stats = compute_episode_stats(episode_buffer, self.features)
- self.meta.save_episode(episode_index, episode_length, task, task_index)
-
- if encode_videos and len(self.meta.video_keys) > 0:
+ if len(self.meta.video_keys) > 0:
video_paths = self.encode_episode_videos(episode_index)
for key in self.meta.video_keys:
episode_buffer[key] = video_paths[key]
+ # `meta.save_episode` be executed after encoding the videos
+ self.meta.save_episode(episode_index, episode_length, episode_tasks, ep_stats)
+
+ ep_data_index = get_episode_data_index(self.meta.episodes, [episode_index])
+ ep_data_index_np = {k: t.numpy() for k, t in ep_data_index.items()}
+ check_timestamps_sync(
+ episode_buffer["timestamp"],
+ episode_buffer["episode_index"],
+ ep_data_index_np,
+ self.fps,
+ self.tolerance_s,
+ )
+
+ video_files = list(self.root.rglob("*.mp4"))
+ assert len(video_files) == self.num_episodes * len(self.meta.video_keys)
+
+ parquet_files = list(self.root.rglob("*.parquet"))
+ assert len(parquet_files) == self.num_episodes
+
+ # delete images
+ img_dir = self.root / "images"
+ if img_dir.is_dir():
+ shutil.rmtree(self.root / "images")
+
if not episode_data: # Reset the buffer
self.episode_buffer = self.create_episode_buffer()
- self.consolidated = False
-
def _save_episode_table(self, episode_buffer: dict, episode_index: int) -> None:
episode_dict = {key: episode_buffer[key] for key in self.hf_features}
ep_dataset = datasets.Dataset.from_dict(episode_dict, features=self.hf_features, split="train")
+ ep_dataset = embed_images(ep_dataset)
+ self.hf_dataset = concatenate_datasets([self.hf_dataset, ep_dataset])
+ self.hf_dataset.set_transform(hf_transform_to_torch)
ep_data_path = self.root / self.meta.get_data_file_path(ep_index=episode_index)
ep_data_path.parent.mkdir(parents=True, exist_ok=True)
- write_parquet(ep_dataset, ep_data_path)
+ ep_dataset.to_parquet(ep_data_path)
def clear_episode_buffer(self) -> None:
episode_index = self.episode_buffer["episode_index"]
@@ -884,38 +984,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
return video_paths
- def consolidate(self, run_compute_stats: bool = True, keep_image_files: bool = False) -> None:
- self.hf_dataset = self.load_hf_dataset()
- self.episode_data_index = get_episode_data_index(self.meta.episodes, self.episodes)
- check_timestamps_sync(self.hf_dataset, self.episode_data_index, self.fps, self.tolerance_s)
-
- if len(self.meta.video_keys) > 0:
- self.encode_videos()
- self.meta.write_video_info()
-
- if not keep_image_files:
- img_dir = self.root / "images"
- if img_dir.is_dir():
- shutil.rmtree(self.root / "images")
-
- video_files = list(self.root.rglob("*.mp4"))
- assert len(video_files) == self.num_episodes * len(self.meta.video_keys)
-
- parquet_files = list(self.root.rglob("*.parquet"))
- assert len(parquet_files) == self.num_episodes
-
- if run_compute_stats:
- self.stop_image_writer()
- # TODO(aliberts): refactor stats in save_episodes
- self.meta.stats = compute_stats(self)
- serialized_stats = serialize_dict(self.meta.stats)
- write_json(serialized_stats, self.root / STATS_PATH)
- self.consolidated = True
- else:
- logging.warning(
- "Skipping computation of the dataset statistics, dataset is not fully consolidated."
- )
-
@classmethod
def create(
cls,
@@ -944,7 +1012,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
)
obj.repo_id = obj.meta.repo_id
obj.root = obj.meta.root
- obj.local_files_only = obj.meta.local_files_only
+ obj.revision = None
obj.tolerance_s = tolerance_s
obj.image_writer = None
@@ -954,19 +1022,13 @@ class LeRobotDataset(torch.utils.data.Dataset):
# TODO(aliberts, rcadene, alexander-soare): Merge this with OnlineBuffer/DataBuffer
obj.episode_buffer = obj.create_episode_buffer()
- # This bool indicates that the current LeRobotDataset instance is in sync with the files on disk. It
- # is used to know when certain operations are need (for instance, computing dataset statistics). In
- # order to be able to push the dataset to the hub, it needs to be consolidated first by calling
- # self.consolidate().
- obj.consolidated = True
-
obj.episodes = None
- obj.hf_dataset = None
+ obj.hf_dataset = obj.create_hf_dataset()
obj.image_transforms = None
obj.delta_timestamps = None
obj.delta_indices = None
obj.episode_data_index = None
- obj.video_backend = video_backend if video_backend is not None else "pyav"
+ obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec()
return obj
@@ -986,12 +1048,11 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
delta_timestamps: dict[list[float]] | None = None,
tolerances_s: dict | None = None,
download_videos: bool = True,
- local_files_only: bool = False,
video_backend: str | None = None,
):
super().__init__()
self.repo_ids = repo_ids
- self.root = Path(root) if root else LEROBOT_HOME
+ self.root = Path(root) if root else HF_LEROBOT_HOME
self.tolerances_s = tolerances_s if tolerances_s else {repo_id: 1e-4 for repo_id in repo_ids}
# Construct the underlying datasets passing everything but `transform` and `delta_timestamps` which
# are handled by this class.
@@ -1004,7 +1065,6 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
delta_timestamps=delta_timestamps,
tolerance_s=self.tolerances_s[repo_id],
download_videos=download_videos,
- local_files_only=local_files_only,
video_backend=video_backend,
)
for repo_id in repo_ids
@@ -1032,7 +1092,10 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
self.image_transforms = image_transforms
self.delta_timestamps = delta_timestamps
- self.stats = aggregate_stats(self._datasets)
+ # TODO(rcadene, aliberts): We should not perform this aggregation for datasets
+ # with multiple robots of different ranges. Instead we should have one normalization
+ # per robot.
+ self.stats = aggregate_stats([dataset.meta.stats for dataset in self._datasets])
@property
def repo_id_to_index(self):
diff --git a/lerobot/common/datasets/push_dataset_to_hub/CODEBASE_VERSION.md b/lerobot/common/datasets/push_dataset_to_hub/CODEBASE_VERSION.md
deleted file mode 100644
index 8fcc8bbe..00000000
--- a/lerobot/common/datasets/push_dataset_to_hub/CODEBASE_VERSION.md
+++ /dev/null
@@ -1,56 +0,0 @@
-## Using / Updating `CODEBASE_VERSION` (for maintainers)
-
-Since our dataset pushed to the hub are decoupled with the evolution of this repo, we ensure compatibility of
-the datasets with our code, we use a `CODEBASE_VERSION` (defined in
-lerobot/common/datasets/lerobot_dataset.py) variable.
-
-For instance, [`lerobot/pusht`](https://huggingface.co/datasets/lerobot/pusht) has many versions to maintain backward compatibility between LeRobot codebase versions:
-- [v1.0](https://huggingface.co/datasets/lerobot/pusht/tree/v1.0)
-- [v1.1](https://huggingface.co/datasets/lerobot/pusht/tree/v1.1)
-- [v1.2](https://huggingface.co/datasets/lerobot/pusht/tree/v1.2)
-- [v1.3](https://huggingface.co/datasets/lerobot/pusht/tree/v1.3)
-- [v1.4](https://huggingface.co/datasets/lerobot/pusht/tree/v1.4)
-- [v1.5](https://huggingface.co/datasets/lerobot/pusht/tree/v1.5)
-- [v1.6](https://huggingface.co/datasets/lerobot/pusht/tree/v1.6) <-- last version
-- [main](https://huggingface.co/datasets/lerobot/pusht/tree/main) <-- points to the last version
-
-Starting with v1.6, every dataset pushed to the hub or saved locally also have this version number in their
-`info.json` metadata.
-
-### Uploading a new dataset
-If you are pushing a new dataset, you don't need to worry about any of the instructions below, nor to be
-compatible with previous codebase versions. The `push_dataset_to_hub.py` script will automatically tag your
-dataset with the current `CODEBASE_VERSION`.
-
-### Updating an existing dataset
-If you want to update an existing dataset, you need to change the `CODEBASE_VERSION` from `lerobot_dataset.py`
-before running `push_dataset_to_hub.py`. This is especially useful if you introduce a breaking change
-intentionally or not (i.e. something not backward compatible such as modifying the reward functions used,
-deleting some frames at the end of an episode, etc.). That way, people running a previous version of the
-codebase won't be affected by your change and backward compatibility is maintained.
-
-However, you will need to update the version of ALL the other datasets so that they have the new
-`CODEBASE_VERSION` as a branch in their hugging face dataset repository. Don't worry, there is an easy way
-that doesn't require to run `push_dataset_to_hub.py`. You can just "branch-out" from the `main` branch on HF
-dataset repo by running this script which corresponds to a `git checkout -b` (so no copy or upload needed):
-
-```python
-from huggingface_hub import HfApi
-
-from lerobot import available_datasets
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION
-
-api = HfApi()
-
-for repo_id in available_datasets:
- dataset_info = api.list_repo_refs(repo_id, repo_type="dataset")
- branches = [b.name for b in dataset_info.branches]
- if CODEBASE_VERSION in branches:
- print(f"{repo_id} already @{CODEBASE_VERSION}, skipping.")
- continue
- else:
- # Now create a branch named after the new version by branching out from "main"
- # which is expected to be the preceding version
- api.create_branch(repo_id, repo_type="dataset", branch=CODEBASE_VERSION, revision="main")
- print(f"{repo_id} successfully updated @{CODEBASE_VERSION}")
-```
diff --git a/lerobot/common/datasets/push_dataset_to_hub/_download_raw.py b/lerobot/common/datasets/push_dataset_to_hub/_download_raw.py
index edeaf093..cc291cea 100644
--- a/lerobot/common/datasets/push_dataset_to_hub/_download_raw.py
+++ b/lerobot/common/datasets/push_dataset_to_hub/_download_raw.py
@@ -152,7 +152,7 @@ def download_raw(raw_dir: Path, repo_id: str):
stacklevel=1,
)
- # Send warning if raw_dir isn't well formated
+ # Send warning if raw_dir isn't well formatted
if raw_dir.parts[-2] != user_id or raw_dir.parts[-1] != dataset_id:
warnings.warn(
f"""`raw_dir` ({raw_dir}) doesn't contain a community or user id `/` the name of the dataset that
diff --git a/lerobot/common/datasets/push_dataset_to_hub/dora_parquet_format.py b/lerobot/common/datasets/push_dataset_to_hub/dora_parquet_format.py
index 95f9c007..acf820bf 100644
--- a/lerobot/common/datasets/push_dataset_to_hub/dora_parquet_format.py
+++ b/lerobot/common/datasets/push_dataset_to_hub/dora_parquet_format.py
@@ -68,11 +68,11 @@ def load_from_raw(raw_dir: Path, videos_dir: Path, fps: int, video: bool, episod
modality_df,
on="timestamp_utc",
# "nearest" is the best option over "backward", since the latter can desynchronizes camera timestamps by
- # matching timestamps that are too far appart, in order to fit the backward constraints. It's not the case for "nearest".
+ # matching timestamps that are too far apart, in order to fit the backward constraints. It's not the case for "nearest".
# However, note that "nearest" might synchronize the reference camera with other cameras on slightly future timestamps.
- # are too far appart.
+ # are too far apart.
direction="nearest",
- tolerance=pd.Timedelta(f"{1/fps} seconds"),
+ tolerance=pd.Timedelta(f"{1 / fps} seconds"),
)
# Remove rows with episode_index -1 which indicates data that correspond to in-between episodes
df = df[df["episode_index"] != -1]
@@ -126,7 +126,7 @@ def load_from_raw(raw_dir: Path, videos_dir: Path, fps: int, video: bool, episod
videos_dir.parent.mkdir(parents=True, exist_ok=True)
videos_dir.symlink_to((raw_dir / "videos").absolute())
- # sanity check the video paths are well formated
+ # sanity check the video paths are well formatted
for key in df:
if "observation.images." not in key:
continue
@@ -143,7 +143,7 @@ def load_from_raw(raw_dir: Path, videos_dir: Path, fps: int, video: bool, episod
# it is the case for video_frame dictionary = [{"path": ..., "timestamp": ...}]
data_dict[key] = [video_frame[0] for video_frame in df[key].values]
- # sanity check the video path is well formated
+ # sanity check the video path is well formatted
video_path = videos_dir.parent / data_dict[key][0]["path"]
if not video_path.exists():
raise ValueError(f"Video file not found in {video_path}")
diff --git a/lerobot/common/datasets/push_dataset_to_hub/openx_rlds_format.py b/lerobot/common/datasets/push_dataset_to_hub/openx_rlds_format.py
index 1f8a5d14..2ffb8369 100644
--- a/lerobot/common/datasets/push_dataset_to_hub/openx_rlds_format.py
+++ b/lerobot/common/datasets/push_dataset_to_hub/openx_rlds_format.py
@@ -17,7 +17,7 @@
For all datasets in the RLDS format.
For https://github.com/google-deepmind/open_x_embodiment (OPENX) datasets.
-NOTE: You need to install tensorflow and tensorflow_datsets before running this script.
+NOTE: You need to install tensorflow and tensorflow_datasets before running this script.
Example:
python lerobot/scripts/push_dataset_to_hub.py \
diff --git a/lerobot/common/datasets/utils.py b/lerobot/common/datasets/utils.py
index 612bac39..7e297b35 100644
--- a/lerobot/common/datasets/utils.py
+++ b/lerobot/common/datasets/utils.py
@@ -13,10 +13,10 @@
# 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 contextlib
import importlib.resources
import json
import logging
-import textwrap
from collections.abc import Iterator
from itertools import accumulate
from pathlib import Path
@@ -27,14 +27,21 @@ from typing import Any
import datasets
import jsonlines
import numpy as np
-import pyarrow.compute as pc
+import packaging.version
import torch
from datasets.table import embed_table_storage
from huggingface_hub import DatasetCard, DatasetCardData, HfApi
+from huggingface_hub.errors import RevisionNotFoundError
from PIL import Image as PILImage
from torchvision import transforms
+from lerobot.common.datasets.backward_compatibility import (
+ V21_MESSAGE,
+ BackwardCompatibilityError,
+ ForwardCompatibilityError,
+)
from lerobot.common.robot_devices.robots.utils import Robot
+from lerobot.common.utils.utils import is_valid_numpy_dtype_string
from lerobot.configs.types import DictLike, FeatureType, PolicyFeature
DEFAULT_CHUNK_SIZE = 1000 # Max number of episodes per chunk
@@ -42,6 +49,7 @@ DEFAULT_CHUNK_SIZE = 1000 # Max number of episodes per chunk
INFO_PATH = "meta/info.json"
EPISODES_PATH = "meta/episodes.jsonl"
STATS_PATH = "meta/stats.json"
+EPISODES_STATS_PATH = "meta/episodes_stats.jsonl"
TASKS_PATH = "meta/tasks.jsonl"
DEFAULT_VIDEO_PATH = "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4"
@@ -112,17 +120,26 @@ def get_nested_item(obj: DictLike, flattened_key: str, sep: str = "/") -> Any:
def serialize_dict(stats: dict[str, torch.Tensor | np.ndarray | dict]) -> dict:
- serialized_dict = {key: value.tolist() for key, value in flatten_dict(stats).items()}
+ serialized_dict = {}
+ for key, value in flatten_dict(stats).items():
+ if isinstance(value, (torch.Tensor, np.ndarray)):
+ serialized_dict[key] = value.tolist()
+ elif isinstance(value, np.generic):
+ serialized_dict[key] = value.item()
+ elif isinstance(value, (int, float)):
+ serialized_dict[key] = value
+ else:
+ raise NotImplementedError(f"The value '{value}' of type '{type(value)}' is not supported.")
return unflatten_dict(serialized_dict)
-def write_parquet(dataset: datasets.Dataset, fpath: Path) -> None:
+def embed_images(dataset: datasets.Dataset) -> datasets.Dataset:
# Embed image bytes into the table before saving to parquet
format = dataset.format
dataset = dataset.with_format("arrow")
dataset = dataset.map(embed_table_storage, batched=False)
dataset = dataset.with_format(**format)
- dataset.to_parquet(fpath)
+ return dataset
def load_json(fpath: Path) -> Any:
@@ -153,6 +170,10 @@ def append_jsonlines(data: dict, fpath: Path) -> None:
writer.write(data)
+def write_info(info: dict, local_dir: Path):
+ write_json(info, local_dir / INFO_PATH)
+
+
def load_info(local_dir: Path) -> dict:
info = load_json(local_dir / INFO_PATH)
for ft in info["features"].values():
@@ -160,29 +181,76 @@ def load_info(local_dir: Path) -> dict:
return info
-def load_stats(local_dir: Path) -> dict:
- if not (local_dir / STATS_PATH).exists():
- return None
- stats = load_json(local_dir / STATS_PATH)
- stats = {key: torch.tensor(value) for key, value in flatten_dict(stats).items()}
+def write_stats(stats: dict, local_dir: Path):
+ serialized_stats = serialize_dict(stats)
+ write_json(serialized_stats, local_dir / STATS_PATH)
+
+
+def cast_stats_to_numpy(stats) -> dict[str, dict[str, np.ndarray]]:
+ stats = {key: np.array(value) for key, value in flatten_dict(stats).items()}
return unflatten_dict(stats)
-def load_tasks(local_dir: Path) -> dict:
+def load_stats(local_dir: Path) -> dict[str, dict[str, np.ndarray]]:
+ if not (local_dir / STATS_PATH).exists():
+ return None
+ stats = load_json(local_dir / STATS_PATH)
+ return cast_stats_to_numpy(stats)
+
+
+def write_task(task_index: int, task: dict, local_dir: Path):
+ task_dict = {
+ "task_index": task_index,
+ "task": task,
+ }
+ append_jsonlines(task_dict, local_dir / TASKS_PATH)
+
+
+def load_tasks(local_dir: Path) -> tuple[dict, dict]:
tasks = load_jsonlines(local_dir / TASKS_PATH)
- return {item["task_index"]: item["task"] for item in sorted(tasks, key=lambda x: x["task_index"])}
+ tasks = {item["task_index"]: item["task"] for item in sorted(tasks, key=lambda x: x["task_index"])}
+ task_to_task_index = {task: task_index for task_index, task in tasks.items()}
+ return tasks, task_to_task_index
+
+
+def write_episode(episode: dict, local_dir: Path):
+ append_jsonlines(episode, local_dir / EPISODES_PATH)
def load_episodes(local_dir: Path) -> dict:
- return load_jsonlines(local_dir / EPISODES_PATH)
+ episodes = load_jsonlines(local_dir / EPISODES_PATH)
+ return {item["episode_index"]: item for item in sorted(episodes, key=lambda x: x["episode_index"])}
-def load_image_as_numpy(fpath: str | Path, dtype="float32", channel_first: bool = True) -> np.ndarray:
+def write_episode_stats(episode_index: int, episode_stats: dict, local_dir: Path):
+ # We wrap episode_stats in a dictionary since `episode_stats["episode_index"]`
+ # is a dictionary of stats and not an integer.
+ episode_stats = {"episode_index": episode_index, "stats": serialize_dict(episode_stats)}
+ append_jsonlines(episode_stats, local_dir / EPISODES_STATS_PATH)
+
+
+def load_episodes_stats(local_dir: Path) -> dict:
+ episodes_stats = load_jsonlines(local_dir / EPISODES_STATS_PATH)
+ return {
+ item["episode_index"]: cast_stats_to_numpy(item["stats"])
+ for item in sorted(episodes_stats, key=lambda x: x["episode_index"])
+ }
+
+
+def backward_compatible_episodes_stats(
+ stats: dict[str, dict[str, np.ndarray]], episodes: list[int]
+) -> dict[str, dict[str, np.ndarray]]:
+ return {ep_idx: stats for ep_idx in episodes}
+
+
+def load_image_as_numpy(
+ fpath: str | Path, dtype: np.dtype = np.float32, channel_first: bool = True
+) -> np.ndarray:
img = PILImage.open(fpath).convert("RGB")
img_array = np.array(img, dtype=dtype)
if channel_first: # (H, W, C) -> (C, H, W)
img_array = np.transpose(img_array, (2, 0, 1))
- if "float" in dtype:
+ if np.issubdtype(dtype, np.floating):
img_array /= 255.0
return img_array
@@ -201,77 +269,95 @@ def hf_transform_to_torch(items_dict: dict[torch.Tensor | None]):
elif first_item is None:
pass
else:
- items_dict[key] = [torch.tensor(x) for x in items_dict[key]]
+ items_dict[key] = [x if isinstance(x, str) else torch.tensor(x) for x in items_dict[key]]
return items_dict
-def _get_major_minor(version: str) -> tuple[int]:
- split = version.strip("v").split(".")
- return int(split[0]), int(split[1])
-
-
-class BackwardCompatibilityError(Exception):
- def __init__(self, repo_id, version):
- message = textwrap.dedent(f"""
- BackwardCompatibilityError: The dataset you requested ({repo_id}) is in {version} format.
-
- We introduced a new format since v2.0 which is not backward compatible with v1.x.
- Please, use our conversion script. Modify the following command with your own task description:
- ```
- python lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py \\
- --repo-id {repo_id} \\
- --single-task "TASK DESCRIPTION." # <---- /!\\ Replace TASK DESCRIPTION /!\\
- ```
-
- A few examples to replace TASK DESCRIPTION: "Pick up the blue cube and place it into the bin.",
- "Insert the peg into the socket.", "Slide open the ziploc bag.", "Take the elevator to the 1st floor.",
- "Open the top cabinet, store the pot inside it then close the cabinet.", "Push the T-shaped block onto the T-shaped target.",
- "Grab the spray paint on the shelf and place it in the bin on top of the robot dog.", "Fold the sweatshirt.", ...
-
- If you encounter a problem, contact LeRobot maintainers on [Discord](https://discord.com/invite/s3KuuzsPFb)
- or open an [issue on GitHub](https://github.com/huggingface/lerobot/issues/new/choose).
- """)
- super().__init__(message)
+def is_valid_version(version: str) -> bool:
+ try:
+ packaging.version.parse(version)
+ return True
+ except packaging.version.InvalidVersion:
+ return False
def check_version_compatibility(
- repo_id: str, version_to_check: str, current_version: str, enforce_breaking_major: bool = True
+ repo_id: str,
+ version_to_check: str | packaging.version.Version,
+ current_version: str | packaging.version.Version,
+ enforce_breaking_major: bool = True,
) -> None:
- current_major, _ = _get_major_minor(current_version)
- major_to_check, _ = _get_major_minor(version_to_check)
- if major_to_check < current_major and enforce_breaking_major:
- raise BackwardCompatibilityError(repo_id, version_to_check)
- elif float(version_to_check.strip("v")) < float(current_version.strip("v")):
- logging.warning(
- f"""The dataset you requested ({repo_id}) was created with a previous version ({version_to_check}) of the
- codebase. The current codebase version is {current_version}. You should be fine since
- backward compatibility is maintained. If you encounter a problem, contact LeRobot maintainers on
- Discord ('https://discord.com/invite/s3KuuzsPFb') or open an issue on github.""",
- )
+ v_check = (
+ packaging.version.parse(version_to_check)
+ if not isinstance(version_to_check, packaging.version.Version)
+ else version_to_check
+ )
+ v_current = (
+ packaging.version.parse(current_version)
+ if not isinstance(current_version, packaging.version.Version)
+ else current_version
+ )
+ if v_check.major < v_current.major and enforce_breaking_major:
+ raise BackwardCompatibilityError(repo_id, v_check)
+ elif v_check.minor < v_current.minor:
+ logging.warning(V21_MESSAGE.format(repo_id=repo_id, version=v_check))
-def get_hub_safe_version(repo_id: str, version: str) -> str:
+def get_repo_versions(repo_id: str) -> list[packaging.version.Version]:
+ """Returns available valid versions (branches and tags) on given repo."""
api = HfApi()
- dataset_info = api.list_repo_refs(repo_id, repo_type="dataset")
- branches = [b.name for b in dataset_info.branches]
- if version not in branches:
- num_version = float(version.strip("v"))
- hub_num_versions = [float(v.strip("v")) for v in branches if v.startswith("v")]
- if num_version >= 2.0 and all(v < 2.0 for v in hub_num_versions):
- raise BackwardCompatibilityError(repo_id, version)
+ repo_refs = api.list_repo_refs(repo_id, repo_type="dataset")
+ repo_refs = [b.name for b in repo_refs.branches + repo_refs.tags]
+ repo_versions = []
+ for ref in repo_refs:
+ with contextlib.suppress(packaging.version.InvalidVersion):
+ repo_versions.append(packaging.version.parse(ref))
- logging.warning(
- f"""You are trying to load a dataset from {repo_id} created with a previous version of the
- codebase. The following versions are available: {branches}.
- The requested version ('{version}') is not found. You should be fine since
- backward compatibility is maintained. If you encounter a problem, contact LeRobot maintainers on
- Discord ('https://discord.com/invite/s3KuuzsPFb') or open an issue on github.""",
+ return repo_versions
+
+
+def get_safe_version(repo_id: str, version: str | packaging.version.Version) -> str:
+ """
+ Returns the version if available on repo or the latest compatible one.
+ Otherwise, will throw a `CompatibilityError`.
+ """
+ target_version = (
+ packaging.version.parse(version) if not isinstance(version, packaging.version.Version) else version
+ )
+ hub_versions = get_repo_versions(repo_id)
+
+ if not hub_versions:
+ raise RevisionNotFoundError(
+ f"""Your dataset must be tagged with a codebase version.
+ Assuming _version_ is the codebase_version value in the info.json, you can run this:
+ ```python
+ from huggingface_hub import HfApi
+
+ hub_api = HfApi()
+ hub_api.create_tag("{repo_id}", tag="_version_", repo_type="dataset")
+ ```
+ """
)
- if "main" not in branches:
- raise ValueError(f"Version 'main' not found on {repo_id}")
- return "main"
- else:
- return version
+
+ if target_version in hub_versions:
+ return f"v{target_version}"
+
+ compatibles = [
+ v for v in hub_versions if v.major == target_version.major and v.minor <= target_version.minor
+ ]
+ if compatibles:
+ return_version = max(compatibles)
+ if return_version < target_version:
+ logging.warning(f"Revision {version} for {repo_id} not found, using version v{return_version}")
+ return f"v{return_version}"
+
+ lower_major = [v for v in hub_versions if v.major < target_version.major]
+ if lower_major:
+ raise BackwardCompatibilityError(repo_id, max(lower_major))
+
+ upper_versions = [v for v in hub_versions if v > target_version]
+ assert len(upper_versions) > 0
+ raise ForwardCompatibilityError(repo_id, min(upper_versions))
def get_hf_features_from_features(features: dict) -> datasets.Features:
@@ -283,11 +369,20 @@ def get_hf_features_from_features(features: dict) -> datasets.Features:
hf_features[key] = datasets.Image()
elif ft["shape"] == (1,):
hf_features[key] = datasets.Value(dtype=ft["dtype"])
- else:
- assert len(ft["shape"]) == 1
+ elif len(ft["shape"]) == 1:
hf_features[key] = datasets.Sequence(
length=ft["shape"][0], feature=datasets.Value(dtype=ft["dtype"])
)
+ elif len(ft["shape"]) == 2:
+ hf_features[key] = datasets.Array2D(shape=ft["shape"], dtype=ft["dtype"])
+ elif len(ft["shape"]) == 3:
+ hf_features[key] = datasets.Array3D(shape=ft["shape"], dtype=ft["dtype"])
+ elif len(ft["shape"]) == 4:
+ hf_features[key] = datasets.Array4D(shape=ft["shape"], dtype=ft["dtype"])
+ elif len(ft["shape"]) == 5:
+ hf_features[key] = datasets.Array5D(shape=ft["shape"], dtype=ft["dtype"])
+ else:
+ raise ValueError(f"Corresponding feature is not valid: {ft}")
return datasets.Features(hf_features)
@@ -358,88 +453,85 @@ def create_empty_dataset_info(
def get_episode_data_index(
- episode_dicts: list[dict], episodes: list[int] | None = None
+ episode_dicts: dict[dict], episodes: list[int] | None = None
) -> dict[str, torch.Tensor]:
- episode_lengths = {ep_idx: ep_dict["length"] for ep_idx, ep_dict in enumerate(episode_dicts)}
+ episode_lengths = {ep_idx: ep_dict["length"] for ep_idx, ep_dict in episode_dicts.items()}
if episodes is not None:
episode_lengths = {ep_idx: episode_lengths[ep_idx] for ep_idx in episodes}
- cumulative_lenghts = list(accumulate(episode_lengths.values()))
+ cumulative_lengths = list(accumulate(episode_lengths.values()))
return {
- "from": torch.LongTensor([0] + cumulative_lenghts[:-1]),
- "to": torch.LongTensor(cumulative_lenghts),
- }
-
-
-def calculate_total_episode(
- hf_dataset: datasets.Dataset, raise_if_not_contiguous: bool = True
-) -> dict[str, torch.Tensor]:
- episode_indices = sorted(hf_dataset.unique("episode_index"))
- total_episodes = len(episode_indices)
- if raise_if_not_contiguous and episode_indices != list(range(total_episodes)):
- raise ValueError("episode_index values are not sorted and contiguous.")
- return total_episodes
-
-
-def calculate_episode_data_index(hf_dataset: datasets.Dataset) -> dict[str, torch.Tensor]:
- episode_lengths = []
- table = hf_dataset.data.table
- total_episodes = calculate_total_episode(hf_dataset)
- for ep_idx in range(total_episodes):
- ep_table = table.filter(pc.equal(table["episode_index"], ep_idx))
- episode_lengths.insert(ep_idx, len(ep_table))
-
- cumulative_lenghts = list(accumulate(episode_lengths))
- return {
- "from": torch.LongTensor([0] + cumulative_lenghts[:-1]),
- "to": torch.LongTensor(cumulative_lenghts),
+ "from": torch.LongTensor([0] + cumulative_lengths[:-1]),
+ "to": torch.LongTensor(cumulative_lengths),
}
def check_timestamps_sync(
- hf_dataset: datasets.Dataset,
- episode_data_index: dict[str, torch.Tensor],
+ timestamps: np.ndarray,
+ episode_indices: np.ndarray,
+ episode_data_index: dict[str, np.ndarray],
fps: int,
tolerance_s: float,
raise_value_error: bool = True,
) -> bool:
"""
- This check is to make sure that each timestamps is separated to the next by 1/fps +/- tolerance to
- account for possible numerical error.
- """
- timestamps = torch.stack(hf_dataset["timestamp"])
- diffs = torch.diff(timestamps)
- within_tolerance = torch.abs(diffs - 1 / fps) <= tolerance_s
+ This check is to make sure that each timestamp is separated from the next by (1/fps) +/- tolerance
+ to account for possible numerical error.
- # We mask differences between the timestamp at the end of an episode
- # and the one at the start of the next episode since these are expected
- # to be outside tolerance.
- mask = torch.ones(len(diffs), dtype=torch.bool)
- ignored_diffs = episode_data_index["to"][:-1] - 1
+ Args:
+ timestamps (np.ndarray): Array of timestamps in seconds.
+ episode_indices (np.ndarray): Array indicating the episode index for each timestamp.
+ episode_data_index (dict[str, np.ndarray]): A dictionary that includes 'to',
+ which identifies indices for the end of each episode.
+ fps (int): Frames per second. Used to check the expected difference between consecutive timestamps.
+ tolerance_s (float): Allowed deviation from the expected (1/fps) difference.
+ raise_value_error (bool): Whether to raise a ValueError if the check fails.
+
+ Returns:
+ bool: True if all checked timestamp differences lie within tolerance, False otherwise.
+
+ Raises:
+ ValueError: If the check fails and `raise_value_error` is True.
+ """
+ if timestamps.shape != episode_indices.shape:
+ raise ValueError(
+ "timestamps and episode_indices should have the same shape. "
+ f"Found {timestamps.shape=} and {episode_indices.shape=}."
+ )
+
+ # Consecutive differences
+ diffs = np.diff(timestamps)
+ within_tolerance = np.abs(diffs - (1.0 / fps)) <= tolerance_s
+
+ # Mask to ignore differences at the boundaries between episodes
+ mask = np.ones(len(diffs), dtype=bool)
+ ignored_diffs = episode_data_index["to"][:-1] - 1 # indices at the end of each episode
mask[ignored_diffs] = False
filtered_within_tolerance = within_tolerance[mask]
- if not torch.all(filtered_within_tolerance):
+ # Check if all remaining diffs are within tolerance
+ if not np.all(filtered_within_tolerance):
# Track original indices before masking
- original_indices = torch.arange(len(diffs))
+ original_indices = np.arange(len(diffs))
filtered_indices = original_indices[mask]
- outside_tolerance_filtered_indices = torch.nonzero(~filtered_within_tolerance) # .squeeze()
+ outside_tolerance_filtered_indices = np.nonzero(~filtered_within_tolerance)[0]
outside_tolerance_indices = filtered_indices[outside_tolerance_filtered_indices]
- episode_indices = torch.stack(hf_dataset["episode_index"])
outside_tolerances = []
for idx in outside_tolerance_indices:
entry = {
"timestamps": [timestamps[idx], timestamps[idx + 1]],
"diff": diffs[idx],
- "episode_index": episode_indices[idx].item(),
+ "episode_index": episode_indices[idx].item()
+ if hasattr(episode_indices[idx], "item")
+ else episode_indices[idx],
}
outside_tolerances.append(entry)
if raise_value_error:
raise ValueError(
f"""One or several timestamps unexpectedly violate the tolerance inside episode range.
- This might be due to synchronization issues with timestamps during data collection.
+ This might be due to synchronization issues during data collection.
\n{pformat(outside_tolerances)}"""
)
return False
@@ -604,3 +696,118 @@ class IterableNamespace(SimpleNamespace):
def keys(self):
return vars(self).keys()
+
+
+def validate_frame(frame: dict, features: dict):
+ optional_features = {"timestamp"}
+ expected_features = (set(features) - set(DEFAULT_FEATURES.keys())) | {"task"}
+ actual_features = set(frame.keys())
+
+ error_message = validate_features_presence(actual_features, expected_features, optional_features)
+
+ if "task" in frame:
+ error_message += validate_feature_string("task", frame["task"])
+
+ common_features = actual_features & (expected_features | optional_features)
+ for name in common_features - {"task"}:
+ error_message += validate_feature_dtype_and_shape(name, features[name], frame[name])
+
+ if error_message:
+ raise ValueError(error_message)
+
+
+def validate_features_presence(
+ actual_features: set[str], expected_features: set[str], optional_features: set[str]
+):
+ error_message = ""
+ missing_features = expected_features - actual_features
+ extra_features = actual_features - (expected_features | optional_features)
+
+ if missing_features or extra_features:
+ error_message += "Feature mismatch in `frame` dictionary:\n"
+ if missing_features:
+ error_message += f"Missing features: {missing_features}\n"
+ if extra_features:
+ error_message += f"Extra features: {extra_features}\n"
+
+ return error_message
+
+
+def validate_feature_dtype_and_shape(name: str, feature: dict, value: np.ndarray | PILImage.Image | str):
+ expected_dtype = feature["dtype"]
+ expected_shape = feature["shape"]
+ if is_valid_numpy_dtype_string(expected_dtype):
+ return validate_feature_numpy_array(name, expected_dtype, expected_shape, value)
+ elif expected_dtype in ["image", "video"]:
+ return validate_feature_image_or_video(name, expected_shape, value)
+ elif expected_dtype == "string":
+ return validate_feature_string(name, value)
+ else:
+ raise NotImplementedError(f"The feature dtype '{expected_dtype}' is not implemented yet.")
+
+
+def validate_feature_numpy_array(
+ name: str, expected_dtype: str, expected_shape: list[int], value: np.ndarray
+):
+ error_message = ""
+ if isinstance(value, np.ndarray):
+ actual_dtype = value.dtype
+ actual_shape = value.shape
+
+ if actual_dtype != np.dtype(expected_dtype):
+ error_message += f"The feature '{name}' of dtype '{actual_dtype}' is not of the expected dtype '{expected_dtype}'.\n"
+
+ if actual_shape != expected_shape:
+ error_message += f"The feature '{name}' of shape '{actual_shape}' does not have the expected shape '{expected_shape}'.\n"
+ else:
+ error_message += f"The feature '{name}' is not a 'np.ndarray'. Expected type is '{expected_dtype}', but type '{type(value)}' provided instead.\n"
+
+ return error_message
+
+
+def validate_feature_image_or_video(name: str, expected_shape: list[str], value: np.ndarray | PILImage.Image):
+ # Note: The check of pixels range ([0,1] for float and [0,255] for uint8) is done by the image writer threads.
+ error_message = ""
+ if isinstance(value, np.ndarray):
+ actual_shape = value.shape
+ c, h, w = expected_shape
+ if len(actual_shape) != 3 or (actual_shape != (c, h, w) and actual_shape != (h, w, c)):
+ error_message += f"The feature '{name}' of shape '{actual_shape}' does not have the expected shape '{(c, h, w)}' or '{(h, w, c)}'.\n"
+ elif isinstance(value, PILImage.Image):
+ pass
+ else:
+ error_message += f"The feature '{name}' is expected to be of type 'PIL.Image' or 'np.ndarray' channel first or channel last, but type '{type(value)}' provided instead.\n"
+
+ return error_message
+
+
+def validate_feature_string(name: str, value: str):
+ if not isinstance(value, str):
+ return f"The feature '{name}' is expected to be of type 'str', but type '{type(value)}' provided instead.\n"
+ return ""
+
+
+def validate_episode_buffer(episode_buffer: dict, total_episodes: int, features: dict):
+ if "size" not in episode_buffer:
+ raise ValueError("size key not found in episode_buffer")
+
+ if "task" not in episode_buffer:
+ raise ValueError("task key not found in episode_buffer")
+
+ if episode_buffer["episode_index"] != total_episodes:
+ # TODO(aliberts): Add option to use existing episode_index
+ raise NotImplementedError(
+ "You might have manually provided the episode_buffer with an episode_index that doesn't "
+ "match the total number of episodes already in the dataset. This is not supported for now."
+ )
+
+ if episode_buffer["size"] == 0:
+ raise ValueError("You must add one or several frames with `add_frame` before calling `add_episode`.")
+
+ buffer_keys = set(episode_buffer.keys()) - {"task", "size"}
+ if not buffer_keys == set(features):
+ raise ValueError(
+ f"Features from `episode_buffer` don't match the ones in `features`."
+ f"In episode_buffer not in features: {buffer_keys - set(features)}"
+ f"In features not in episode_buffer: {set(features) - buffer_keys}"
+ )
diff --git a/lerobot/common/datasets/v2/batch_convert_dataset_v1_to_v2.py b/lerobot/common/datasets/v2/batch_convert_dataset_v1_to_v2.py
index 4cd93a2d..99ab2cbf 100644
--- a/lerobot/common/datasets/v2/batch_convert_dataset_v1_to_v2.py
+++ b/lerobot/common/datasets/v2/batch_convert_dataset_v1_to_v2.py
@@ -31,6 +31,7 @@ from lerobot.common.robot_devices.robots.configs import AlohaRobotConfig
LOCAL_DIR = Path("data/")
+# spellchecker:off
ALOHA_MOBILE_INFO = {
"robot_config": AlohaRobotConfig(),
"license": "mit",
@@ -856,6 +857,7 @@ DATASETS = {
}""").lstrip(),
},
}
+# spellchecker:on
def batch_convert():
diff --git a/lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py b/lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py
index 62ca9932..acf0282f 100644
--- a/lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py
+++ b/lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py
@@ -17,7 +17,7 @@
"""
This script will help you convert any LeRobot dataset already pushed to the hub from codebase version 1.6 to
2.0. You will be required to provide the 'tasks', which is a short but accurate description in plain English
-for each of the task performed in the dataset. This will allow to easily train models with task-conditionning.
+for each of the task performed in the dataset. This will allow to easily train models with task-conditioning.
We support 3 different scenarios for these tasks (see instructions below):
1. Single task dataset: all episodes of your dataset have the same single task.
@@ -130,7 +130,7 @@ from lerobot.common.datasets.utils import (
create_branch,
create_lerobot_dataset_card,
flatten_dict,
- get_hub_safe_version,
+ get_safe_version,
load_json,
unflatten_dict,
write_json,
@@ -443,7 +443,7 @@ def convert_dataset(
test_branch: str | None = None,
**card_kwargs,
):
- v1 = get_hub_safe_version(repo_id, V16)
+ v1 = get_safe_version(repo_id, V16)
v1x_dir = local_dir / V16 / repo_id
v20_dir = local_dir / V20 / repo_id
v1x_dir.mkdir(parents=True, exist_ok=True)
diff --git a/lerobot/common/datasets/v21/_remove_language_instruction.py b/lerobot/common/datasets/v21/_remove_language_instruction.py
new file mode 100644
index 00000000..643ddd3f
--- /dev/null
+++ b/lerobot/common/datasets/v21/_remove_language_instruction.py
@@ -0,0 +1,87 @@
+# 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 logging
+import traceback
+from pathlib import Path
+
+from datasets import get_dataset_config_info
+from huggingface_hub import HfApi
+
+from lerobot import available_datasets
+from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
+from lerobot.common.datasets.utils import INFO_PATH, write_info
+from lerobot.common.datasets.v21.convert_dataset_v20_to_v21 import V20, SuppressWarnings
+
+LOCAL_DIR = Path("data/")
+
+hub_api = HfApi()
+
+
+def fix_dataset(repo_id: str) -> str:
+ if not hub_api.revision_exists(repo_id, V20, repo_type="dataset"):
+ return f"{repo_id}: skipped (not in {V20})."
+
+ dataset_info = get_dataset_config_info(repo_id, "default")
+ with SuppressWarnings():
+ lerobot_metadata = LeRobotDatasetMetadata(repo_id, revision=V20, force_cache_sync=True)
+
+ meta_features = {key for key, ft in lerobot_metadata.features.items() if ft["dtype"] != "video"}
+ parquet_features = set(dataset_info.features)
+
+ diff_parquet_meta = parquet_features - meta_features
+ diff_meta_parquet = meta_features - parquet_features
+
+ if diff_parquet_meta:
+ raise ValueError(f"In parquet not in info.json: {parquet_features - meta_features}")
+
+ if not diff_meta_parquet:
+ return f"{repo_id}: skipped (no diff)"
+
+ if diff_meta_parquet:
+ logging.warning(f"In info.json not in parquet: {meta_features - parquet_features}")
+ assert diff_meta_parquet == {"language_instruction"}
+ lerobot_metadata.features.pop("language_instruction")
+ write_info(lerobot_metadata.info, lerobot_metadata.root)
+ commit_info = hub_api.upload_file(
+ path_or_fileobj=lerobot_metadata.root / INFO_PATH,
+ path_in_repo=INFO_PATH,
+ repo_id=repo_id,
+ repo_type="dataset",
+ revision=V20,
+ commit_message="Remove 'language_instruction'",
+ create_pr=True,
+ )
+ return f"{repo_id}: success - PR: {commit_info.pr_url}"
+
+
+def batch_fix():
+ status = {}
+ LOCAL_DIR.mkdir(parents=True, exist_ok=True)
+ logfile = LOCAL_DIR / "fix_features_v20.txt"
+ for num, repo_id in enumerate(available_datasets):
+ print(f"\nConverting {repo_id} ({num}/{len(available_datasets)})")
+ print("---------------------------------------------------------")
+ try:
+ status = fix_dataset(repo_id)
+ except Exception:
+ status = f"{repo_id}: failed\n {traceback.format_exc()}"
+
+ logging.info(status)
+ with open(logfile, "a") as file:
+ file.write(status + "\n")
+
+
+if __name__ == "__main__":
+ batch_fix()
diff --git a/lerobot/common/datasets/v21/batch_convert_dataset_v20_to_v21.py b/lerobot/common/datasets/v21/batch_convert_dataset_v20_to_v21.py
new file mode 100644
index 00000000..cc9272a8
--- /dev/null
+++ b/lerobot/common/datasets/v21/batch_convert_dataset_v20_to_v21.py
@@ -0,0 +1,54 @@
+#!/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.
+
+"""
+This script is for internal use to convert all datasets under the 'lerobot' hub user account to v2.1.
+"""
+
+import traceback
+from pathlib import Path
+
+from huggingface_hub import HfApi
+
+from lerobot import available_datasets
+from lerobot.common.datasets.v21.convert_dataset_v20_to_v21 import V21, convert_dataset
+
+LOCAL_DIR = Path("data/")
+
+
+def batch_convert():
+ status = {}
+ LOCAL_DIR.mkdir(parents=True, exist_ok=True)
+ logfile = LOCAL_DIR / "conversion_log_v21.txt"
+ hub_api = HfApi()
+ for num, repo_id in enumerate(available_datasets):
+ print(f"\nConverting {repo_id} ({num}/{len(available_datasets)})")
+ print("---------------------------------------------------------")
+ try:
+ if hub_api.revision_exists(repo_id, V21, repo_type="dataset"):
+ status = f"{repo_id}: success (already in {V21})."
+ else:
+ convert_dataset(repo_id)
+ status = f"{repo_id}: success."
+ except Exception:
+ status = f"{repo_id}: failed\n {traceback.format_exc()}"
+
+ with open(logfile, "a") as file:
+ file.write(status + "\n")
+
+
+if __name__ == "__main__":
+ batch_convert()
diff --git a/lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py b/lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py
new file mode 100644
index 00000000..176d16d0
--- /dev/null
+++ b/lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py
@@ -0,0 +1,114 @@
+# 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.
+
+"""
+This script will help you convert any LeRobot dataset already pushed to the hub from codebase version 2.0 to
+2.1. 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/v21/convert_dataset_v20_to_v21.py \
+ --repo-id=aliberts/koch_tutorial
+```
+
+"""
+
+import argparse
+import logging
+
+from huggingface_hub import HfApi
+
+from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
+from lerobot.common.datasets.utils import EPISODES_STATS_PATH, STATS_PATH, load_stats, write_info
+from lerobot.common.datasets.v21.convert_stats import check_aggregate_stats, convert_stats
+
+V20 = "v2.0"
+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,
+):
+ with SuppressWarnings():
+ dataset = LeRobotDataset(repo_id, revision=V20, force_cache_sync=True)
+
+ if (dataset.root / EPISODES_STATS_PATH).is_file():
+ (dataset.root / EPISODES_STATS_PATH).unlink()
+
+ convert_stats(dataset, num_workers=num_workers)
+ ref_stats = load_stats(dataset.root)
+ check_aggregate_stats(dataset, ref_stats)
+
+ dataset.meta.info["codebase_version"] = CODEBASE_VERSION
+ write_info(dataset.meta.info, dataset.root)
+
+ dataset.push_to_hub(branch=branch, tag_version=False, allow_patterns="meta/")
+
+ # delete old stats.json file
+ if (dataset.root / STATS_PATH).is_file:
+ (dataset.root / STATS_PATH).unlink()
+
+ hub_api = HfApi()
+ if hub_api.file_exists(
+ repo_id=dataset.repo_id, filename=STATS_PATH, revision=branch, repo_type="dataset"
+ ):
+ hub_api.delete_file(
+ path_in_repo=STATS_PATH, repo_id=dataset.repo_id, revision=branch, repo_type="dataset"
+ )
+
+ hub_api.create_tag(repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
+
+
+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))
diff --git a/lerobot/common/datasets/v21/convert_stats.py b/lerobot/common/datasets/v21/convert_stats.py
new file mode 100644
index 00000000..4a20b427
--- /dev/null
+++ b/lerobot/common/datasets/v21/convert_stats.py
@@ -0,0 +1,99 @@
+# 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.
+
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+import numpy as np
+from tqdm import tqdm
+
+from lerobot.common.datasets.compute_stats import aggregate_stats, get_feature_stats, sample_indices
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.common.datasets.utils import write_episode_stats
+
+
+def sample_episode_video_frames(dataset: LeRobotDataset, episode_index: int, ft_key: str) -> np.ndarray:
+ ep_len = dataset.meta.episodes[episode_index]["length"]
+ sampled_indices = sample_indices(ep_len)
+ query_timestamps = dataset._get_query_timestamps(0.0, {ft_key: sampled_indices})
+ video_frames = dataset._query_videos(query_timestamps, episode_index)
+ return video_frames[ft_key].numpy()
+
+
+def convert_episode_stats(dataset: LeRobotDataset, ep_idx: int):
+ ep_start_idx = dataset.episode_data_index["from"][ep_idx]
+ ep_end_idx = dataset.episode_data_index["to"][ep_idx]
+ ep_data = dataset.hf_dataset.select(range(ep_start_idx, ep_end_idx))
+
+ ep_stats = {}
+ for key, ft in dataset.features.items():
+ if ft["dtype"] == "video":
+ # We sample only for videos
+ ep_ft_data = sample_episode_video_frames(dataset, ep_idx, key)
+ else:
+ ep_ft_data = np.array(ep_data[key])
+
+ axes_to_reduce = (0, 2, 3) if ft["dtype"] in ["image", "video"] else 0
+ keepdims = True if ft["dtype"] in ["image", "video"] else ep_ft_data.ndim == 1
+ ep_stats[key] = get_feature_stats(ep_ft_data, axis=axes_to_reduce, keepdims=keepdims)
+
+ if ft["dtype"] in ["image", "video"]: # remove batch dim
+ ep_stats[key] = {
+ k: v if k == "count" else np.squeeze(v, axis=0) for k, v in ep_stats[key].items()
+ }
+
+ dataset.meta.episodes_stats[ep_idx] = ep_stats
+
+
+def convert_stats(dataset: LeRobotDataset, num_workers: int = 0):
+ assert dataset.episodes is None
+ print("Computing episodes stats")
+ total_episodes = dataset.meta.total_episodes
+ if num_workers > 0:
+ with ThreadPoolExecutor(max_workers=num_workers) as executor:
+ futures = {
+ executor.submit(convert_episode_stats, dataset, ep_idx): ep_idx
+ for ep_idx in range(total_episodes)
+ }
+ for future in tqdm(as_completed(futures), total=total_episodes):
+ future.result()
+ else:
+ for ep_idx in tqdm(range(total_episodes)):
+ convert_episode_stats(dataset, ep_idx)
+
+ for ep_idx in tqdm(range(total_episodes)):
+ write_episode_stats(ep_idx, dataset.meta.episodes_stats[ep_idx], dataset.root)
+
+
+def check_aggregate_stats(
+ dataset: LeRobotDataset,
+ reference_stats: dict[str, dict[str, np.ndarray]],
+ video_rtol_atol: tuple[float] = (1e-2, 1e-2),
+ default_rtol_atol: tuple[float] = (5e-6, 6e-5),
+):
+ """Verifies that the aggregated stats from episodes_stats are close to reference stats."""
+ agg_stats = aggregate_stats(list(dataset.meta.episodes_stats.values()))
+ for key, ft in dataset.features.items():
+ # These values might need some fine-tuning
+ if ft["dtype"] == "video":
+ # to account for image sub-sampling
+ rtol, atol = video_rtol_atol
+ else:
+ rtol, atol = default_rtol_atol
+
+ for stat, val in agg_stats[key].items():
+ if key in reference_stats and stat in reference_stats[key]:
+ err_msg = f"feature='{key}' stats='{stat}'"
+ np.testing.assert_allclose(
+ val, reference_stats[key][stat], rtol=rtol, atol=atol, err_msg=err_msg
+ )
diff --git a/lerobot/common/datasets/video_utils.py b/lerobot/common/datasets/video_utils.py
index 8ed3318d..4f696861 100644
--- a/lerobot/common/datasets/video_utils.py
+++ b/lerobot/common/datasets/video_utils.py
@@ -13,6 +13,7 @@
# 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 json
import logging
import subprocess
@@ -29,6 +30,46 @@ from datasets.features.features import register_feature
from PIL import Image
+def get_safe_default_codec():
+ if importlib.util.find_spec("torchcodec"):
+ return "torchcodec"
+ else:
+ logging.warning(
+ "'torchcodec' is not available in your platform, falling back to 'pyav' as a default decoder"
+ )
+ return "pyav"
+
+
+def decode_video_frames(
+ video_path: Path | str,
+ timestamps: list[float],
+ tolerance_s: float,
+ backend: str | None = None,
+) -> torch.Tensor:
+ """
+ Decodes video frames using the specified backend.
+
+ Args:
+ video_path (Path): Path to the video file.
+ timestamps (list[float]): List of timestamps to extract frames.
+ tolerance_s (float): Allowed deviation in seconds for frame retrieval.
+ backend (str, optional): Backend to use for decoding. Defaults to "torchcodec" when available in the platform; otherwise, defaults to "pyav"..
+
+ Returns:
+ torch.Tensor: Decoded frames.
+
+ Currently supports torchcodec on cpu and pyav.
+ """
+ if backend is None:
+ backend = get_safe_default_codec()
+ if backend == "torchcodec":
+ return decode_video_frames_torchcodec(video_path, timestamps, tolerance_s)
+ elif backend in ["pyav", "video_reader"]:
+ return decode_video_frames_torchvision(video_path, timestamps, tolerance_s, backend)
+ else:
+ raise ValueError(f"Unsupported video backend: {backend}")
+
+
def decode_video_frames_torchvision(
video_path: Path | str,
timestamps: list[float],
@@ -69,11 +110,11 @@ def decode_video_frames_torchvision(
# set the first and last requested timestamps
# Note: previous timestamps are usually loaded, since we need to access the previous key frame
- first_ts = timestamps[0]
- last_ts = timestamps[-1]
+ first_ts = min(timestamps)
+ last_ts = max(timestamps)
# access closest key frame of the first requested frame
- # Note: closest key frame timestamp is usally smaller than `first_ts` (e.g. key frame can be the first frame of the video)
+ # Note: closest key frame timestamp is usually smaller than `first_ts` (e.g. key frame can be the first frame of the video)
# for details on what `seek` is doing see: https://pyav.basswood-io.com/docs/stable/api/container.html?highlight=inputcontainer#av.container.InputContainer.seek
reader.seek(first_ts, keyframes_only=keyframes_only)
@@ -127,6 +168,81 @@ def decode_video_frames_torchvision(
return closest_frames
+def decode_video_frames_torchcodec(
+ video_path: Path | str,
+ timestamps: list[float],
+ tolerance_s: float,
+ device: str = "cpu",
+ log_loaded_timestamps: bool = False,
+) -> torch.Tensor:
+ """Loads frames associated with the requested timestamps of a video using torchcodec.
+
+ Note: Setting device="cuda" outside the main process, e.g. in data loader workers, will lead to CUDA initialization errors.
+
+ Note: Video benefits from inter-frame compression. Instead of storing every frame individually,
+ the encoder stores a reference frame (or a key frame) and subsequent frames as differences relative to
+ that key frame. As a consequence, to access a requested frame, we need to load the preceding key frame,
+ and all subsequent frames until reaching the requested frame. The number of key frames in a video
+ can be adjusted during encoding to take into account decoding time and video size in bytes.
+ """
+
+ if importlib.util.find_spec("torchcodec"):
+ from torchcodec.decoders import VideoDecoder
+ else:
+ raise ImportError("torchcodec is required but not available.")
+
+ # initialize video decoder
+ decoder = VideoDecoder(video_path, device=device, seek_mode="approximate")
+ loaded_frames = []
+ loaded_ts = []
+ # get metadata for frame information
+ metadata = decoder.metadata
+ average_fps = metadata.average_fps
+
+ # convert timestamps to frame indices
+ frame_indices = [round(ts * average_fps) for ts in timestamps]
+
+ # retrieve frames based on indices
+ frames_batch = decoder.get_frames_at(indices=frame_indices)
+
+ for frame, pts in zip(frames_batch.data, frames_batch.pts_seconds, strict=False):
+ loaded_frames.append(frame)
+ loaded_ts.append(pts.item())
+ if log_loaded_timestamps:
+ logging.info(f"Frame loaded at timestamp={pts:.4f}")
+
+ query_ts = torch.tensor(timestamps)
+ loaded_ts = torch.tensor(loaded_ts)
+
+ # compute distances between each query timestamp and loaded timestamps
+ dist = torch.cdist(query_ts[:, None], loaded_ts[:, None], p=1)
+ min_, argmin_ = dist.min(1)
+
+ is_within_tol = min_ < tolerance_s
+ assert is_within_tol.all(), (
+ f"One or several query timestamps unexpectedly violate the tolerance ({min_[~is_within_tol]} > {tolerance_s=})."
+ "It means that the closest frame that can be loaded from the video is too far away in time."
+ "This might be due to synchronization issues with timestamps during data collection."
+ "To be safe, we advise to ignore this item during training."
+ f"\nqueried timestamps: {query_ts}"
+ f"\nloaded timestamps: {loaded_ts}"
+ f"\nvideo: {video_path}"
+ )
+
+ # get closest frames to the query timestamps
+ closest_frames = torch.stack([loaded_frames[idx] for idx in argmin_])
+ closest_ts = loaded_ts[argmin_]
+
+ if log_loaded_timestamps:
+ logging.info(f"{closest_ts=}")
+
+ # convert to float32 in [0,1] range (channel first)
+ closest_frames = closest_frames.type(torch.float32) / 255
+
+ assert len(timestamps) == len(closest_frames)
+ return closest_frames
+
+
def encode_video_frames(
imgs_dir: Path | str,
video_path: Path | str,
diff --git a/lerobot/common/envs/__init__.py b/lerobot/common/envs/__init__.py
index a583ffc5..4977d11d 100644
--- a/lerobot/common/envs/__init__.py
+++ b/lerobot/common/envs/__init__.py
@@ -1 +1,15 @@
+# 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.
+
from .configs import AlohaEnv, EnvConfig, PushtEnv, XarmEnv # noqa: F401
diff --git a/lerobot/common/envs/configs.py b/lerobot/common/envs/configs.py
index 6259ca94..cf90048a 100644
--- a/lerobot/common/envs/configs.py
+++ b/lerobot/common/envs/configs.py
@@ -1,3 +1,17 @@
+# 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 abc
from dataclasses import dataclass, field
diff --git a/lerobot/common/envs/factory.py b/lerobot/common/envs/factory.py
index 49239363..8450f84b 100644
--- a/lerobot/common/envs/factory.py
+++ b/lerobot/common/envs/factory.py
@@ -37,12 +37,12 @@ def make_env(cfg: EnvConfig, n_envs: int = 1, use_async_envs: bool = False) -> g
Args:
cfg (EnvConfig): the config of the environment to instantiate.
n_envs (int, optional): The number of parallelized env to return. Defaults to 1.
- use_async_envs (bool, optional): Wether to return an AsyncVectorEnv or a SyncVectorEnv. Defaults to
+ use_async_envs (bool, optional): Whether to return an AsyncVectorEnv or a SyncVectorEnv. Defaults to
False.
Raises:
ValueError: if n_envs < 1
- ModuleNotFoundError: If the requested env package is not intalled
+ ModuleNotFoundError: If the requested env package is not installed
Returns:
gym.vector.VectorEnv: The parallelized gym.env instance.
diff --git a/lerobot/common/logger.py b/lerobot/common/logger.py
deleted file mode 100644
index 5f863f68..00000000
--- a/lerobot/common/logger.py
+++ /dev/null
@@ -1,240 +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.
-"""Borrowed from https://github.com/fyhMer/fowm/blob/main/src/logger.py
-
-# TODO(rcadene, alexander-soare): clean this file
-"""
-
-import logging
-import os
-import re
-from dataclasses import asdict
-from glob import glob
-from pathlib import Path
-
-import draccus
-import torch
-from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
-from termcolor import colored
-from torch.optim import Optimizer
-from torch.optim.lr_scheduler import LRScheduler
-
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.utils.utils import get_global_random_state
-from lerobot.configs.train import TrainPipelineConfig
-from lerobot.configs.types import FeatureType, NormalizationMode
-
-PRETRAINED_MODEL = "pretrained_model"
-TRAINING_STATE = "training_state.pth"
-
-
-def log_output_dir(out_dir):
- logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {out_dir}")
-
-
-def cfg_to_group(cfg: TrainPipelineConfig, return_list: bool = False) -> list[str] | str:
- """Return a group name for logging. Optionally returns group name as list."""
- lst = [
- f"policy:{cfg.policy.type}",
- f"dataset:{cfg.dataset.repo_id}",
- f"seed:{cfg.seed}",
- ]
- if cfg.env is not None:
- lst.append(f"env:{cfg.env.type}")
- return lst if return_list else "-".join(lst)
-
-
-def get_wandb_run_id_from_filesystem(checkpoint_dir: Path) -> str:
- # Get the WandB run ID.
- paths = glob(str(checkpoint_dir / "../wandb/latest-run/run-*"))
- if len(paths) != 1:
- raise RuntimeError("Couldn't get the previous WandB run ID for run resumption.")
- match = re.search(r"run-([^\.]+).wandb", paths[0].split("/")[-1])
- if match is None:
- raise RuntimeError("Couldn't get the previous WandB run ID for run resumption.")
- wandb_run_id = match.groups(0)[0]
- return wandb_run_id
-
-
-class Logger:
- """Primary logger object. Logs either locally or using wandb.
-
- The logger creates the following directory structure:
-
- provided_log_dir
- ├── checkpoints
- │ ├── specific_checkpoint_name
- │ │ ├── pretrained_model # Hugging Face pretrained model directory
- │ │ │ ├── ...
- │ │ └── training_state.pth # optimizer, scheduler, and random states + training step
- | ├── another_specific_checkpoint_name
- │ │ ├── ...
- | ├── ...
- │ └── last # a softlink to the last logged checkpoint
- """
-
- pretrained_model_dir_name = PRETRAINED_MODEL
- training_state_file_name = TRAINING_STATE
-
- def __init__(self, cfg: TrainPipelineConfig):
- self._cfg = cfg
- self.log_dir = cfg.output_dir
- self.log_dir.mkdir(parents=True, exist_ok=True)
- self.job_name = cfg.job_name
- self.checkpoints_dir = self.get_checkpoints_dir(self.log_dir)
- self.last_checkpoint_dir = self.get_last_checkpoint_dir(self.log_dir)
- self.last_pretrained_model_dir = self.get_last_pretrained_model_dir(self.log_dir)
-
- # Set up WandB.
- self._group = cfg_to_group(cfg)
- run_offline = not cfg.wandb.enable or not cfg.wandb.project
- if run_offline:
- logging.info(colored("Logs will be saved locally.", "yellow", attrs=["bold"]))
- self._wandb = None
- else:
- os.environ["WANDB_SILENT"] = "true"
- import wandb
-
- wandb_run_id = None
- if cfg.resume:
- wandb_run_id = get_wandb_run_id_from_filesystem(self.checkpoints_dir)
-
- wandb.init(
- id=wandb_run_id,
- project=cfg.wandb.project,
- entity=cfg.wandb.entity,
- name=self.job_name,
- notes=cfg.wandb.notes,
- tags=cfg_to_group(cfg, return_list=True),
- dir=self.log_dir,
- config=asdict(self._cfg),
- # TODO(rcadene): try set to True
- save_code=False,
- # TODO(rcadene): split train and eval, and run async eval with job_type="eval"
- job_type="train_eval",
- resume="must" if cfg.resume else None,
- )
- print(colored("Logs will be synced with wandb.", "blue", attrs=["bold"]))
- logging.info(f"Track this run --> {colored(wandb.run.get_url(), 'yellow', attrs=['bold'])}")
- self._wandb = wandb
-
- @classmethod
- def get_checkpoints_dir(cls, log_dir: str | Path) -> Path:
- """Given the log directory, get the sub-directory in which checkpoints will be saved."""
- return Path(log_dir) / "checkpoints"
-
- @classmethod
- def get_last_checkpoint_dir(cls, log_dir: str | Path) -> Path:
- """Given the log directory, get the sub-directory in which the last checkpoint will be saved."""
- return cls.get_checkpoints_dir(log_dir) / "last"
-
- @classmethod
- def get_last_pretrained_model_dir(cls, log_dir: str | Path) -> Path:
- """
- Given the log directory, get the sub-directory in which the last checkpoint's pretrained weights will
- be saved.
- """
- return cls.get_last_checkpoint_dir(log_dir) / cls.pretrained_model_dir_name
-
- def save_model(self, save_dir: Path, policy: PreTrainedPolicy, wandb_artifact_name: str | None = None):
- """Save the weights of the Policy model using PyTorchModelHubMixin.
-
- The weights are saved in a folder called "pretrained_model" under the checkpoint directory.
-
- Optionally also upload the model to WandB.
- """
-
- self.checkpoints_dir.mkdir(parents=True, exist_ok=True)
- register_features_types()
- policy.save_pretrained(save_dir)
- # Also save the full config for the env configuration.
- self._cfg.save_pretrained(save_dir)
- if self._wandb and not self._cfg.wandb.disable_artifact:
- # note wandb artifact does not accept ":" or "/" in its name
- artifact = self._wandb.Artifact(wandb_artifact_name, type="model")
- artifact.add_file(save_dir / SAFETENSORS_SINGLE_FILE)
- self._wandb.log_artifact(artifact)
- if self.last_checkpoint_dir.exists():
- os.remove(self.last_checkpoint_dir)
-
- def save_training_state(
- self,
- save_dir: Path,
- train_step: int,
- optimizer: Optimizer | None = None,
- scheduler: LRScheduler | None = None,
- ):
- """Checkpoint the global training_step, optimizer state, scheduler state, and random state.
-
- All of these are saved as "training_state.pth" under the checkpoint directory.
- """
- training_state = {}
- training_state["step"] = train_step
- training_state.update(get_global_random_state())
- if optimizer is not None:
- training_state["optimizer"] = optimizer.state_dict()
- if scheduler is not None:
- training_state["scheduler"] = scheduler.state_dict()
- torch.save(training_state, save_dir / self.training_state_file_name)
-
- def save_checkpoint(
- self,
- train_step: int,
- identifier: str,
- policy: PreTrainedPolicy,
- optimizer: Optimizer | None = None,
- scheduler: LRScheduler | None = None,
- ):
- """Checkpoint the model weights and the training state."""
- checkpoint_dir = self.checkpoints_dir / str(identifier)
- wandb_artifact_name = (
- None
- if self._wandb is None
- else f"{self._group.replace(':', '_').replace('/', '_')}-{self._cfg.seed}-{identifier}"
- )
- self.save_model(
- checkpoint_dir / self.pretrained_model_dir_name, policy, wandb_artifact_name=wandb_artifact_name
- )
- self.save_training_state(checkpoint_dir, train_step, optimizer, scheduler)
-
- relative_target = checkpoint_dir.relative_to(self.last_checkpoint_dir.parent)
- self.last_checkpoint_dir.symlink_to(relative_target)
-
- def log_dict(self, d: dict, step: int, mode: str = "train"):
- assert mode in {"train", "eval"}
- # TODO(alexander-soare): Add local text log.
- if self._wandb is not None:
- for k, v in d.items():
- if not isinstance(v, (int, float, str)):
- logging.warning(
- f'WandB logging of key "{k}" was ignored as its type is not handled by this wrapper.'
- )
- continue
- self._wandb.log({f"{mode}/{k}": v}, step=step)
-
- def log_video(self, video_path: str, step: int, mode: str = "train"):
- assert mode in {"train", "eval"}
- assert self._wandb is not None
- wandb_video = self._wandb.Video(video_path, fps=self._cfg.env.fps, format="mp4")
- self._wandb.log({f"{mode}/video": wandb_video}, step=step)
-
-
-def register_features_types():
- draccus.decode.register(FeatureType, lambda x: FeatureType[x])
- draccus.encode.register(FeatureType, lambda x: x.name)
-
- draccus.decode.register(NormalizationMode, lambda x: NormalizationMode[x])
- draccus.encode.register(NormalizationMode, lambda x: x.name)
diff --git a/lerobot/common/optim/__init__.py b/lerobot/common/optim/__init__.py
index e1e65966..de2c4c99 100644
--- a/lerobot/common/optim/__init__.py
+++ b/lerobot/common/optim/__init__.py
@@ -1 +1,15 @@
+# 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.
+
from .optimizers import OptimizerConfig as OptimizerConfig
diff --git a/lerobot/common/optim/factory.py b/lerobot/common/optim/factory.py
index 010cd461..10ff3df7 100644
--- a/lerobot/common/optim/factory.py
+++ b/lerobot/common/optim/factory.py
@@ -14,15 +14,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from pathlib import Path
-import torch
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LRScheduler
-from lerobot.common.logger import TRAINING_STATE
from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.utils.utils import get_global_random_state, set_global_random_state
from lerobot.configs.train import TrainPipelineConfig
@@ -40,22 +36,5 @@ def make_optimizer_and_scheduler(
"""
params = policy.get_optim_params() if cfg.use_policy_training_preset else policy.parameters()
optimizer = cfg.optimizer.build(params)
- lr_scheduler = cfg.scheduler.build(optimizer, cfg.offline.steps) if cfg.scheduler is not None else None
+ lr_scheduler = cfg.scheduler.build(optimizer, cfg.steps) if cfg.scheduler is not None else None
return optimizer, lr_scheduler
-
-
-def load_training_state(checkpoint_dir: Path, optimizer: Optimizer, scheduler: LRScheduler | None) -> int:
- """
- Given the checkpoint directory, load the optimizer state, scheduler state, and random state, and
- return the global training step.
- """
- # TODO(aliberts): use safetensors instead as weights_only=False is unsafe
- training_state = torch.load(checkpoint_dir / TRAINING_STATE, weights_only=False)
- optimizer.load_state_dict(training_state["optimizer"])
- if scheduler is not None:
- scheduler.load_state_dict(training_state["scheduler"])
- elif "scheduler" in training_state:
- raise ValueError("The checkpoint contains a scheduler state_dict, but no LRScheduler was provided.")
- # Small HACK to get the expected keys: use `get_global_random_state`.
- set_global_random_state({k: training_state[k] for k in get_global_random_state()})
- return training_state["step"], optimizer, scheduler
diff --git a/lerobot/common/optim/optimizers.py b/lerobot/common/optim/optimizers.py
index 737305ad..0cf4124c 100644
--- a/lerobot/common/optim/optimizers.py
+++ b/lerobot/common/optim/optimizers.py
@@ -1,8 +1,32 @@
+#!/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 abc
from dataclasses import asdict, dataclass
+from pathlib import Path
import draccus
import torch
+from safetensors.torch import load_file, save_file
+
+from lerobot.common.constants import (
+ OPTIMIZER_PARAM_GROUPS,
+ OPTIMIZER_STATE,
+)
+from lerobot.common.datasets.utils import flatten_dict, unflatten_dict, write_json
+from lerobot.common.utils.io_utils import deserialize_json_into_object
@dataclass
@@ -68,3 +92,27 @@ class SGDConfig(OptimizerConfig):
kwargs = asdict(self)
kwargs.pop("grad_clip_norm")
return torch.optim.SGD(params, **kwargs)
+
+
+def save_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> None:
+ state = optimizer.state_dict()
+ param_groups = state.pop("param_groups")
+ flat_state = flatten_dict(state)
+ save_file(flat_state, save_dir / OPTIMIZER_STATE)
+ write_json(param_groups, save_dir / OPTIMIZER_PARAM_GROUPS)
+
+
+def load_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> torch.optim.Optimizer:
+ current_state_dict = optimizer.state_dict()
+ flat_state = load_file(save_dir / OPTIMIZER_STATE)
+ state = unflatten_dict(flat_state)
+ loaded_state_dict = {"state": {int(k): v for k, v in state["state"].items()}}
+
+ if "param_groups" in current_state_dict:
+ param_groups = deserialize_json_into_object(
+ save_dir / OPTIMIZER_PARAM_GROUPS, current_state_dict["param_groups"]
+ )
+ loaded_state_dict["param_groups"] = param_groups
+
+ optimizer.load_state_dict(loaded_state_dict)
+ return optimizer
diff --git a/lerobot/common/optim/schedulers.py b/lerobot/common/optim/schedulers.py
index 80d83bdf..7e158394 100644
--- a/lerobot/common/optim/schedulers.py
+++ b/lerobot/common/optim/schedulers.py
@@ -1,11 +1,31 @@
+#!/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 abc
import math
from dataclasses import asdict, dataclass
+from pathlib import Path
import draccus
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR, LRScheduler
+from lerobot.common.constants import SCHEDULER_STATE
+from lerobot.common.datasets.utils import write_json
+from lerobot.common.utils.io_utils import deserialize_json_into_object
+
@dataclass
class LRSchedulerConfig(draccus.ChoiceRegistry, abc.ABC):
@@ -89,3 +109,14 @@ class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig):
return cosine_decay_schedule(current_step)
return LambdaLR(optimizer, lr_lambda, -1)
+
+
+def save_scheduler_state(scheduler: LRScheduler, save_dir: Path) -> None:
+ state_dict = scheduler.state_dict()
+ write_json(state_dict, save_dir / SCHEDULER_STATE)
+
+
+def load_scheduler_state(scheduler: LRScheduler, save_dir: Path) -> LRScheduler:
+ state_dict = deserialize_json_into_object(save_dir / SCHEDULER_STATE, scheduler.state_dict())
+ scheduler.load_state_dict(state_dict)
+ return scheduler
diff --git a/lerobot/common/policies/__init__.py b/lerobot/common/policies/__init__.py
index 2e4486ef..b73ba5f4 100644
--- a/lerobot/common/policies/__init__.py
+++ b/lerobot/common/policies/__init__.py
@@ -1,3 +1,17 @@
+# 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.
+
from .act.configuration_act import ACTConfig as ACTConfig
from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
from .pi0.configuration_pi0 import PI0Config as PI0Config
diff --git a/lerobot/common/policies/act/configuration_act.py b/lerobot/common/policies/act/configuration_act.py
index 4f724c12..7a5819b7 100644
--- a/lerobot/common/policies/act/configuration_act.py
+++ b/lerobot/common/policies/act/configuration_act.py
@@ -64,7 +64,7 @@ class ACTConfig(PreTrainedConfig):
output_normalization_modes: Similar dictionary as `normalize_input_modes`, but to unnormalize to the
original scale. Note that this is also used for normalizing the training targets.
vision_backbone: Name of the torchvision resnet backbone to use for encoding images.
- pretrained_backbone_weights: Pretrained weights from torchvision to initalize the backbone.
+ pretrained_backbone_weights: Pretrained weights from torchvision to initialize the backbone.
`None` means no pretrained weights.
replace_final_stride_with_dilation: Whether to replace the ResNet's final 2x2 stride with a dilated
convolution.
diff --git a/lerobot/common/policies/act/modeling_act.py b/lerobot/common/policies/act/modeling_act.py
index 615f238f..72d4df03 100644
--- a/lerobot/common/policies/act/modeling_act.py
+++ b/lerobot/common/policies/act/modeling_act.py
@@ -119,9 +119,7 @@ class ACTPolicy(PreTrainedPolicy):
batch = self.normalize_inputs(batch)
if self.config.image_features:
batch = dict(batch) # shallow copy so that adding a key doesn't modify the original
- batch["observation.images"] = torch.stack(
- [batch[key] for key in self.config.image_features], dim=-4
- )
+ batch["observation.images"] = [batch[key] for key in self.config.image_features]
# If we are doing temporal ensembling, do online updates where we keep track of the number of actions
# we are ensembling over.
@@ -144,14 +142,13 @@ class ACTPolicy(PreTrainedPolicy):
self._action_queue.extend(actions.transpose(0, 1))
return self._action_queue.popleft()
- def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+ def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
"""Run the batch through the model and compute the loss for training or validation."""
batch = self.normalize_inputs(batch)
if self.config.image_features:
batch = dict(batch) # shallow copy so that adding a key doesn't modify the original
- batch["observation.images"] = torch.stack(
- [batch[key] for key in self.config.image_features], dim=-4
- )
+ batch["observation.images"] = [batch[key] for key in self.config.image_features]
+
batch = self.normalize_targets(batch)
actions_hat, (mu_hat, log_sigma_x2_hat) = self.model(batch)
@@ -169,11 +166,11 @@ class ACTPolicy(PreTrainedPolicy):
(-0.5 * (1 + log_sigma_x2_hat - mu_hat.pow(2) - (log_sigma_x2_hat).exp())).sum(-1).mean()
)
loss_dict["kld_loss"] = mean_kld.item()
- loss_dict["loss"] = l1_loss + mean_kld * self.config.kl_weight
+ loss = l1_loss + mean_kld * self.config.kl_weight
else:
- loss_dict["loss"] = l1_loss
+ loss = l1_loss
- return loss_dict
+ return loss, loss_dict
class ACTTemporalEnsembler:
@@ -409,15 +406,14 @@ class ACT(nn.Module):
latent dimension.
"""
if self.config.use_vae and self.training:
- assert (
- "action" in batch
- ), "actions must be provided when using the variational objective in training mode."
+ assert "action" in batch, (
+ "actions must be provided when using the variational objective in training mode."
+ )
- batch_size = (
- batch["observation.images"]
- if "observation.images" in batch
- else batch["observation.environment_state"]
- ).shape[0]
+ if "observation.images" in batch:
+ batch_size = batch["observation.images"][0].shape[0]
+ else:
+ batch_size = batch["observation.environment_state"].shape[0]
# Prepare the latent for input to the transformer encoder.
if self.config.use_vae and "action" in batch:
@@ -490,20 +486,21 @@ class ACT(nn.Module):
all_cam_features = []
all_cam_pos_embeds = []
- for cam_index in range(batch["observation.images"].shape[-4]):
- cam_features = self.backbone(batch["observation.images"][:, cam_index])["feature_map"]
- # TODO(rcadene, alexander-soare): remove call to `.to` to speedup forward ; precompute and use
- # buffer
+ # For a list of images, the H and W may vary but H*W is constant.
+ for img in batch["observation.images"]:
+ cam_features = self.backbone(img)["feature_map"]
cam_pos_embed = self.encoder_cam_feat_pos_embed(cam_features).to(dtype=cam_features.dtype)
- cam_features = self.encoder_img_feat_input_proj(cam_features) # (B, C, h, w)
+ cam_features = self.encoder_img_feat_input_proj(cam_features)
+
+ # Rearrange features to (sequence, batch, dim).
+ cam_features = einops.rearrange(cam_features, "b c h w -> (h w) b c")
+ cam_pos_embed = einops.rearrange(cam_pos_embed, "b c h w -> (h w) b c")
+
all_cam_features.append(cam_features)
all_cam_pos_embeds.append(cam_pos_embed)
- # Concatenate camera observation feature maps and positional embeddings along the width dimension,
- # and move to (sequence, batch, dim).
- all_cam_features = torch.cat(all_cam_features, axis=-1)
- encoder_in_tokens.extend(einops.rearrange(all_cam_features, "b c h w -> (h w) b c"))
- all_cam_pos_embeds = torch.cat(all_cam_pos_embeds, axis=-1)
- encoder_in_pos_embed.extend(einops.rearrange(all_cam_pos_embeds, "b c h w -> (h w) b c"))
+
+ encoder_in_tokens.extend(torch.cat(all_cam_features, axis=0))
+ encoder_in_pos_embed.extend(torch.cat(all_cam_pos_embeds, axis=0))
# Stack all tokens along the sequence dimension.
encoder_in_tokens = torch.stack(encoder_in_tokens, axis=0)
diff --git a/lerobot/common/policies/diffusion/configuration_diffusion.py b/lerobot/common/policies/diffusion/configuration_diffusion.py
index 31d5dc8b..e73c65fe 100644
--- a/lerobot/common/policies/diffusion/configuration_diffusion.py
+++ b/lerobot/common/policies/diffusion/configuration_diffusion.py
@@ -68,7 +68,7 @@ class DiffusionConfig(PreTrainedConfig):
within the image size. If None, no cropping is done.
crop_is_random: Whether the crop should be random at training time (it's always a center crop in eval
mode).
- pretrained_backbone_weights: Pretrained weights from torchvision to initalize the backbone.
+ pretrained_backbone_weights: Pretrained weights from torchvision to initialize the backbone.
`None` means no pretrained weights.
use_group_norm: Whether to replace batch normalization with group normalization in the backbone.
The group sizes are set to be about 16 (to be precise, feature_dim // 16).
@@ -99,7 +99,7 @@ class DiffusionConfig(PreTrainedConfig):
num_inference_steps: Number of reverse diffusion steps to use at inference time (steps are evenly
spaced). If not provided, this defaults to be the same as `num_train_timesteps`.
do_mask_loss_for_padding: Whether to mask the loss when there are copy-padded actions. See
- `LeRobotDataset` and `load_previous_and_future_frames` for mor information. Note, this defaults
+ `LeRobotDataset` and `load_previous_and_future_frames` for more information. Note, this defaults
to False as the original Diffusion Policy implementation does the same.
"""
@@ -221,7 +221,7 @@ class DiffusionConfig(PreTrainedConfig):
for key, image_ft in self.image_features.items():
if image_ft.shape != first_image_ft.shape:
raise ValueError(
- f"`{key}` does not match `{first_image_key}`, but we " "expect all image shapes to match."
+ f"`{key}` does not match `{first_image_key}`, but we expect all image shapes to match."
)
@property
diff --git a/lerobot/common/policies/diffusion/modeling_diffusion.py b/lerobot/common/policies/diffusion/modeling_diffusion.py
index 7147f550..9ecadcb0 100644
--- a/lerobot/common/policies/diffusion/modeling_diffusion.py
+++ b/lerobot/common/policies/diffusion/modeling_diffusion.py
@@ -143,7 +143,7 @@ class DiffusionPolicy(PreTrainedPolicy):
action = self._queues["action"].popleft()
return action
- def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+ def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, None]:
"""Run the batch through the model and compute the loss for training or validation."""
batch = self.normalize_inputs(batch)
if self.config.image_features:
@@ -153,7 +153,8 @@ class DiffusionPolicy(PreTrainedPolicy):
)
batch = self.normalize_targets(batch)
loss = self.diffusion.compute_loss(batch)
- return {"loss": loss}
+ # no output_dict so returning None
+ return loss, None
def _make_noise_scheduler(name: str, **kwargs: dict) -> DDPMScheduler | DDIMScheduler:
diff --git a/lerobot/common/policies/factory.py b/lerobot/common/policies/factory.py
index cd440f7a..5d2f6cb5 100644
--- a/lerobot/common/policies/factory.py
+++ b/lerobot/common/policies/factory.py
@@ -16,7 +16,6 @@
import logging
-import torch
from torch import nn
from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
@@ -76,7 +75,6 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
def make_policy(
cfg: PreTrainedConfig,
- device: str | torch.device,
ds_meta: LeRobotDatasetMetadata | None = None,
env_cfg: EnvConfig | None = None,
) -> PreTrainedPolicy:
@@ -88,7 +86,6 @@ def make_policy(
Args:
cfg (PreTrainedConfig): The config of the policy to make. If `pretrained_path` is set, the policy will
be loaded with the weights from that path.
- device (str): the device to load the policy onto.
ds_meta (LeRobotDatasetMetadata | None, optional): Dataset metadata to take input/output shapes and
statistics to use for (un)normalization of inputs/outputs in the policy. Defaults to None.
env_cfg (EnvConfig | None, optional): The config of a gym environment to parse features from. Must be
@@ -96,7 +93,7 @@ def make_policy(
Raises:
ValueError: Either ds_meta or env and env_cfg must be provided.
- NotImplementedError: if the policy.type is 'vqbet' and the device 'mps' (due to an incompatibility)
+ NotImplementedError: if the policy.type is 'vqbet' and the policy device 'mps' (due to an incompatibility)
Returns:
PreTrainedPolicy: _description_
@@ -111,7 +108,7 @@ def make_policy(
# https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment
# variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be
# slower than running natively on MPS.
- if cfg.type == "vqbet" and str(device) == "mps":
+ if cfg.type == "vqbet" and cfg.device == "mps":
raise NotImplementedError(
"Current implementation of VQBeT does not support `mps` backend. "
"Please use `cpu` or `cuda` backend."
@@ -145,7 +142,7 @@ def make_policy(
# Make a fresh policy.
policy = policy_cls(**kwargs)
- policy.to(device)
+ policy.to(cfg.device)
assert isinstance(policy, nn.Module)
# policy = torch.compile(policy, mode="reduce-overhead")
diff --git a/lerobot/common/policies/normalize.py b/lerobot/common/policies/normalize.py
index 95219273..b3255ec1 100644
--- a/lerobot/common/policies/normalize.py
+++ b/lerobot/common/policies/normalize.py
@@ -13,6 +13,7 @@
# 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 numpy as np
import torch
from torch import Tensor, nn
@@ -77,17 +78,29 @@ def create_stats_buffers(
}
)
+ # TODO(aliberts, rcadene): harmonize this to only use one framework (np or torch)
if stats:
- # Note: The clone is needed to make sure that the logic in save_pretrained doesn't see duplicated
- # tensors anywhere (for example, when we use the same stats for normalization and
- # unnormalization). See the logic here
- # https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L97.
- if norm_mode is NormalizationMode.MEAN_STD:
- buffer["mean"].data = stats[key]["mean"].clone()
- buffer["std"].data = stats[key]["std"].clone()
- elif norm_mode is NormalizationMode.MIN_MAX:
- buffer["min"].data = stats[key]["min"].clone()
- buffer["max"].data = stats[key]["max"].clone()
+ if isinstance(stats[key]["mean"], np.ndarray):
+ if norm_mode is NormalizationMode.MEAN_STD:
+ buffer["mean"].data = torch.from_numpy(stats[key]["mean"]).to(dtype=torch.float32)
+ buffer["std"].data = torch.from_numpy(stats[key]["std"]).to(dtype=torch.float32)
+ elif norm_mode is NormalizationMode.MIN_MAX:
+ buffer["min"].data = torch.from_numpy(stats[key]["min"]).to(dtype=torch.float32)
+ buffer["max"].data = torch.from_numpy(stats[key]["max"]).to(dtype=torch.float32)
+ elif isinstance(stats[key]["mean"], torch.Tensor):
+ # Note: The clone is needed to make sure that the logic in save_pretrained doesn't see duplicated
+ # tensors anywhere (for example, when we use the same stats for normalization and
+ # unnormalization). See the logic here
+ # https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L97.
+ if norm_mode is NormalizationMode.MEAN_STD:
+ buffer["mean"].data = stats[key]["mean"].clone().to(dtype=torch.float32)
+ buffer["std"].data = stats[key]["std"].clone().to(dtype=torch.float32)
+ elif norm_mode is NormalizationMode.MIN_MAX:
+ buffer["min"].data = stats[key]["min"].clone().to(dtype=torch.float32)
+ buffer["max"].data = stats[key]["max"].clone().to(dtype=torch.float32)
+ else:
+ type_ = type(stats[key]["mean"])
+ raise ValueError(f"np.ndarray or torch.Tensor expected, but type is '{type_}' instead.")
stats_buffers[key] = buffer
return stats_buffers
@@ -141,6 +154,7 @@ class Normalize(nn.Module):
batch = dict(batch) # shallow copy avoids mutating the input batch
for key, ft in self.features.items():
if key not in batch:
+ # FIXME(aliberts, rcadene): This might lead to silent fail!
continue
norm_mode = self.norm_map.get(ft.type, NormalizationMode.IDENTITY)
diff --git a/lerobot/common/policies/pi0/configuration_pi0.py b/lerobot/common/policies/pi0/configuration_pi0.py
index 8d2eedf6..8c7cc130 100644
--- a/lerobot/common/policies/pi0/configuration_pi0.py
+++ b/lerobot/common/policies/pi0/configuration_pi0.py
@@ -1,3 +1,17 @@
+# 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.
+
from dataclasses import dataclass, field
from lerobot.common.optim.optimizers import AdamWConfig
@@ -76,6 +90,7 @@ class PI0Config(PreTrainedConfig):
def __post_init__(self):
super().__post_init__()
+ # TODO(Steven): Validate device and amp? in all policy configs?
"""Input validation (not exhaustive)."""
if self.n_action_steps > self.chunk_size:
raise ValueError(
diff --git a/lerobot/common/policies/pi0/conversion_scripts/benchmark.py b/lerobot/common/policies/pi0/conversion_scripts/benchmark.py
index 31bd1b66..cb3c0e9b 100644
--- a/lerobot/common/policies/pi0/conversion_scripts/benchmark.py
+++ b/lerobot/common/policies/pi0/conversion_scripts/benchmark.py
@@ -1,3 +1,17 @@
+# 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 torch
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
@@ -31,7 +45,7 @@ def main():
cfg = PreTrainedConfig.from_pretrained(ckpt_torch_dir)
cfg.pretrained_path = ckpt_torch_dir
- policy = make_policy(cfg, device, ds_meta=dataset.meta)
+ policy = make_policy(cfg, ds_meta=dataset.meta)
# policy = torch.compile(policy, mode="reduce-overhead")
diff --git a/lerobot/common/policies/pi0/conversion_scripts/compare_with_jax.py b/lerobot/common/policies/pi0/conversion_scripts/compare_with_jax.py
index 8b2e1c66..6bd7c91f 100644
--- a/lerobot/common/policies/pi0/conversion_scripts/compare_with_jax.py
+++ b/lerobot/common/policies/pi0/conversion_scripts/compare_with_jax.py
@@ -1,3 +1,17 @@
+# 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 json
import pickle
from pathlib import Path
@@ -87,7 +101,7 @@ def main():
cfg = PreTrainedConfig.from_pretrained(ckpt_torch_dir)
cfg.pretrained_path = ckpt_torch_dir
- policy = make_policy(cfg, device, dataset_meta)
+ policy = make_policy(cfg, dataset_meta)
# loss_dict = policy.forward(batch, noise=noise, time=time_beta)
# loss_dict["loss"].backward()
diff --git a/lerobot/common/policies/pi0/conversion_scripts/conversion_utils.py b/lerobot/common/policies/pi0/conversion_scripts/conversion_utils.py
index 8e35d0d4..8835da31 100644
--- a/lerobot/common/policies/pi0/conversion_scripts/conversion_utils.py
+++ b/lerobot/common/policies/pi0/conversion_scripts/conversion_utils.py
@@ -1,3 +1,17 @@
+# 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.
+
from transformers import GemmaConfig, PaliGemmaConfig
diff --git a/lerobot/common/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py b/lerobot/common/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py
index f85437a5..73ff506f 100644
--- a/lerobot/common/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py
+++ b/lerobot/common/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py
@@ -1,8 +1,22 @@
+# 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.
+
"""
Convert pi0 parameters from Jax to Pytorch
Follow [README of openpi](https://github.com/Physical-Intelligence/openpi) to create a new environment
-and install the required librairies.
+and install the required libraries.
```bash
cd ~/code/openpi
diff --git a/lerobot/common/policies/pi0/flex_attention.py b/lerobot/common/policies/pi0/flex_attention.py
index 38a5b597..35628cdd 100644
--- a/lerobot/common/policies/pi0/flex_attention.py
+++ b/lerobot/common/policies/pi0/flex_attention.py
@@ -1,3 +1,17 @@
+# 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 torch
import torch.nn.functional as F # noqa: N812
from packaging.version import Version
diff --git a/lerobot/common/policies/pi0/modeling_pi0.py b/lerobot/common/policies/pi0/modeling_pi0.py
index 90d1a14c..bc53bf85 100644
--- a/lerobot/common/policies/pi0/modeling_pi0.py
+++ b/lerobot/common/policies/pi0/modeling_pi0.py
@@ -300,7 +300,7 @@ class PI0Policy(PreTrainedPolicy):
self._action_queue.extend(actions.transpose(0, 1))
return self._action_queue.popleft()
- def forward(self, batch: dict[str, Tensor], noise=None, time=None) -> dict[str, Tensor]:
+ def forward(self, batch: dict[str, Tensor], noise=None, time=None) -> tuple[Tensor, dict[str, Tensor]]:
"""Do a full training forward pass to compute the loss"""
if self.config.adapt_to_pi_aloha:
batch[OBS_ROBOT] = self._pi_aloha_decode_state(batch[OBS_ROBOT])
@@ -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)
@@ -328,12 +328,12 @@ class PI0Policy(PreTrainedPolicy):
losses = losses[:, :, : self.config.max_action_dim]
loss_dict["losses_after_rm_padding"] = losses.clone()
- loss = losses.mean()
# For backward pass
- loss_dict["loss"] = loss
+ loss = losses.mean()
# For logging
loss_dict["l2_loss"] = loss.item()
- return loss_dict
+
+ return loss, loss_dict
def prepare_images(self, batch):
"""Apply Pi0 preprocessing to the images, like resizing to 224x224 and padding to keep aspect ratio, and
diff --git a/lerobot/common/policies/pi0/paligemma_with_expert.py b/lerobot/common/policies/pi0/paligemma_with_expert.py
index 08c36c11..76e2ce60 100644
--- a/lerobot/common/policies/pi0/paligemma_with_expert.py
+++ b/lerobot/common/policies/pi0/paligemma_with_expert.py
@@ -1,3 +1,17 @@
+# 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.
+
from typing import List, Optional, Union
import torch
diff --git a/lerobot/common/policies/pretrained.py b/lerobot/common/policies/pretrained.py
index 84767594..da4ef157 100644
--- a/lerobot/common/policies/pretrained.py
+++ b/lerobot/common/policies/pretrained.py
@@ -1,3 +1,16 @@
+# 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 abc
import logging
import os
@@ -73,7 +86,6 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
cache_dir: str | Path | None = None,
local_files_only: bool = False,
revision: str | None = None,
- map_location: str = "cpu",
strict: bool = False,
**kwargs,
) -> T:
@@ -98,7 +110,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
if os.path.isdir(model_id):
print("Loading weights from local directory")
model_file = os.path.join(model_id, SAFETENSORS_SINGLE_FILE)
- policy = cls._load_as_safetensor(instance, model_file, map_location, strict)
+ policy = cls._load_as_safetensor(instance, model_file, config.device, strict)
else:
try:
model_file = hf_hub_download(
@@ -112,13 +124,13 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
token=token,
local_files_only=local_files_only,
)
- policy = cls._load_as_safetensor(instance, model_file, map_location, strict)
+ policy = cls._load_as_safetensor(instance, model_file, config.device, strict)
except HfHubHTTPError as e:
raise FileNotFoundError(
f"{SAFETENSORS_SINGLE_FILE} not found on the HuggingFace Hub in {model_id}"
) from e
- policy.to(map_location)
+ policy.to(config.device)
policy.eval()
return policy
@@ -163,12 +175,17 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
"""
raise NotImplementedError
+ # TODO(aliberts, rcadene): split into 'forward' and 'compute_loss'?
@abc.abstractmethod
- def forward(self, batch: dict[str, Tensor]) -> dict:
- """Run the batch through the model and compute the loss for training or validation.
+ def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
+ """_summary_
- Returns a dictionary with "loss" and potentially other information. Apart from "loss" which is a Tensor, all
- other items should be logging-friendly, native Python types.
+ Args:
+ batch (dict[str, Tensor]): _description_
+
+ Returns:
+ tuple[Tensor, dict | None]: The loss and potentially other information. Apart from the loss which
+ is a Tensor, all other items should be logging-friendly, native Python types.
"""
raise NotImplementedError
diff --git a/lerobot/common/policies/tdmpc/configuration_tdmpc.py b/lerobot/common/policies/tdmpc/configuration_tdmpc.py
index c3e8aee6..3fce01df 100644
--- a/lerobot/common/policies/tdmpc/configuration_tdmpc.py
+++ b/lerobot/common/policies/tdmpc/configuration_tdmpc.py
@@ -76,7 +76,7 @@ class TDMPCConfig(PreTrainedConfig):
n_pi_samples: Number of samples to draw from the policy / world model rollout every CEM iteration. Can
be zero.
uncertainty_regularizer_coeff: Coefficient for the uncertainty regularization used when estimating
- trajectory values (this is the λ coeffiecient in eqn 4 of FOWM).
+ trajectory values (this is the λ coefficient in eqn 4 of FOWM).
n_elites: The number of elite samples to use for updating the gaussian parameters every CEM iteration.
elite_weighting_temperature: The temperature to use for softmax weighting (by trajectory value) of the
elites, when updating the gaussian parameters for CEM.
@@ -165,7 +165,7 @@ class TDMPCConfig(PreTrainedConfig):
"""Input validation (not exhaustive)."""
if self.n_gaussian_samples <= 0:
raise ValueError(
- f"The number of guassian samples for CEM should be non-zero. Got `{self.n_gaussian_samples=}`"
+ f"The number of gaussian samples for CEM should be non-zero. Got `{self.n_gaussian_samples=}`"
)
if self.normalization_mapping["ACTION"] is not NormalizationMode.MIN_MAX:
raise ValueError(
diff --git a/lerobot/common/policies/tdmpc/modeling_tdmpc.py b/lerobot/common/policies/tdmpc/modeling_tdmpc.py
index 6366a5a4..0940f198 100644
--- a/lerobot/common/policies/tdmpc/modeling_tdmpc.py
+++ b/lerobot/common/policies/tdmpc/modeling_tdmpc.py
@@ -302,7 +302,7 @@ class TDMPCPolicy(PreTrainedPolicy):
G -= running_discount * self.config.uncertainty_regularizer_coeff * terminal_values.std(0)
return G
- def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor | float]:
+ def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
"""Run the batch through the model and compute the loss.
Returns a dictionary with loss as a tensor, and other information as native floats.
@@ -495,7 +495,6 @@ class TDMPCPolicy(PreTrainedPolicy):
"Q_value_loss": q_value_loss.item(),
"V_value_loss": v_value_loss.item(),
"pi_loss": pi_loss.item(),
- "loss": loss,
"sum_loss": loss.item() * self.config.horizon,
}
)
@@ -505,7 +504,7 @@ class TDMPCPolicy(PreTrainedPolicy):
if isinstance(batch[key], torch.Tensor) and batch[key].ndim > 1:
batch[key] = batch[key].transpose(1, 0)
- return info
+ return loss, info
def update(self):
"""Update the target model's parameters with an EMA step."""
@@ -595,9 +594,9 @@ class TDMPCTOLD(nn.Module):
self.apply(_apply_fn)
for m in [self._reward, *self._Qs]:
- assert isinstance(
- m[-1], nn.Linear
- ), "Sanity check. The last linear layer needs 0 initialization on weights."
+ assert isinstance(m[-1], nn.Linear), (
+ "Sanity check. The last linear layer needs 0 initialization on weights."
+ )
nn.init.zeros_(m[-1].weight)
nn.init.zeros_(m[-1].bias) # this has already been done, but keep this line here for good measure
diff --git a/lerobot/common/policies/vqbet/configuration_vqbet.py b/lerobot/common/policies/vqbet/configuration_vqbet.py
index 47007e82..28e9c433 100644
--- a/lerobot/common/policies/vqbet/configuration_vqbet.py
+++ b/lerobot/common/policies/vqbet/configuration_vqbet.py
@@ -66,7 +66,7 @@ class VQBeTConfig(PreTrainedConfig):
within the image size. If None, no cropping is done.
crop_is_random: Whether the crop should be random at training time (it's always a center crop in eval
mode).
- pretrained_backbone_weights: Pretrained weights from torchvision to initalize the backbone.
+ pretrained_backbone_weights: Pretrained weights from torchvision to initialize the backbone.
`None` means no pretrained weights.
use_group_norm: Whether to replace batch normalization with group normalization in the backbone.
The group sizes are set to be about 16 (to be precise, feature_dim // 16).
@@ -184,7 +184,7 @@ class VQBeTConfig(PreTrainedConfig):
for key, image_ft in self.image_features.items():
if image_ft.shape != first_image_ft.shape:
raise ValueError(
- f"`{key}` does not match `{first_image_key}`, but we " "expect all image shapes to match."
+ f"`{key}` does not match `{first_image_key}`, but we expect all image shapes to match."
)
@property
diff --git a/lerobot/common/policies/vqbet/modeling_vqbet.py b/lerobot/common/policies/vqbet/modeling_vqbet.py
index c4d4a46d..97a08e2f 100644
--- a/lerobot/common/policies/vqbet/modeling_vqbet.py
+++ b/lerobot/common/policies/vqbet/modeling_vqbet.py
@@ -156,7 +156,7 @@ class VQBeTPolicy(PreTrainedPolicy):
action = self._queues["action"].popleft()
return action
- def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+ def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
"""Run the batch through the model and compute the loss for training or validation."""
batch = self.normalize_inputs(batch)
batch = dict(batch) # shallow copy so that adding a key doesn't modify the original
@@ -170,16 +170,16 @@ class VQBeTPolicy(PreTrainedPolicy):
loss, n_different_codes, n_different_combinations, recon_l1_error = (
self.vqbet.action_head.discretize(self.config.n_vqvae_training_steps, batch["action"])
)
- return {
- "loss": loss,
+ return loss, {
"n_different_codes": n_different_codes,
"n_different_combinations": n_different_combinations,
"recon_l1_error": recon_l1_error,
}
# if Residual VQ is already trained, VQ-BeT trains its GPT and bin prediction head / offset prediction head parts.
_, loss_dict = self.vqbet(batch, rollout=False)
+ loss = loss_dict.pop("loss")
- return loss_dict
+ return loss, loss_dict
class SpatialSoftmax(nn.Module):
@@ -342,7 +342,7 @@ class VQBeTModel(nn.Module):
torch.row_stack([torch.arange(i, i + self.config.action_chunk_size) for i in range(num_tokens)]),
)
- def forward(self, batch: dict[str, Tensor], rollout: bool) -> Tensor:
+ def forward(self, batch: dict[str, Tensor], rollout: bool) -> tuple[dict, dict]:
# Input validation.
assert set(batch).issuperset({"observation.state", "observation.images"})
batch_size, n_obs_steps = batch["observation.state"].shape[:2]
@@ -482,10 +482,10 @@ class VQBeTHead(nn.Module):
param.requires_grad = False
return loss, n_different_codes, n_different_combinations, recon_l1_error
- def forward(self, x, **kwargs):
+ def forward(self, x, **kwargs) -> dict:
# N is the batch size, and T is number of action query tokens, which are process through same GPT
N, T, _ = x.shape
- # we calculate N and T side parallely. Thus, the dimensions would be
+ # we calculate N and T side parallelly. Thus, the dimensions would be
# (batch size * number of action query tokens, action chunk size, action dimension)
x = einops.rearrange(x, "N T WA -> (N T) WA")
@@ -772,7 +772,7 @@ class VqVae(nn.Module):
Encoder and decoder are MLPs consisting of an input, output layer, and hidden layer, respectively.
The vq_layer uses residual VQs.
- This class contains functions for training the encoder and decoder along with the residual VQ layer (for trainign phase 1),
+ This class contains functions for training the encoder and decoder along with the residual VQ layer (for training phase 1),
as well as functions to help BeT training part in training phase 2.
"""
diff --git a/lerobot/common/policies/vqbet/vqbet_utils.py b/lerobot/common/policies/vqbet/vqbet_utils.py
index 90a2cfda..139d119e 100644
--- a/lerobot/common/policies/vqbet/vqbet_utils.py
+++ b/lerobot/common/policies/vqbet/vqbet_utils.py
@@ -38,7 +38,7 @@ from lerobot.common.policies.vqbet.configuration_vqbet import VQBeTConfig
This file is part of a VQ-BeT that utilizes code from the following repositories:
- Vector Quantize PyTorch code is licensed under the MIT License:
- Origianl source: https://github.com/lucidrains/vector-quantize-pytorch
+ Original source: https://github.com/lucidrains/vector-quantize-pytorch
- nanoGPT part is an adaptation of Andrej Karpathy's nanoGPT implementation in PyTorch.
Original source: https://github.com/karpathy/nanoGPT
@@ -203,9 +203,9 @@ class GPT(nn.Module):
def forward(self, input, targets=None):
device = input.device
b, t, d = input.size()
- assert (
- t <= self.config.gpt_block_size
- ), f"Cannot forward sequence of length {t}, block size is only {self.config.gpt_block_size}"
+ assert t <= self.config.gpt_block_size, (
+ f"Cannot forward sequence of length {t}, block size is only {self.config.gpt_block_size}"
+ )
# positional encodings that are added to the input embeddings
pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0) # shape (1, t)
@@ -273,10 +273,10 @@ class GPT(nn.Module):
assert len(inter_params) == 0, "parameters {} made it into both decay/no_decay sets!".format(
str(inter_params)
)
- assert (
- len(param_dict.keys() - union_params) == 0
- ), "parameters {} were not separated into either decay/no_decay set!".format(
- str(param_dict.keys() - union_params),
+ assert len(param_dict.keys() - union_params) == 0, (
+ "parameters {} were not separated into either decay/no_decay set!".format(
+ str(param_dict.keys() - union_params),
+ )
)
decay = [param_dict[pn] for pn in sorted(decay)]
@@ -289,7 +289,7 @@ class GPT(nn.Module):
This file is a part for Residual Vector Quantization that utilizes code from the following repository:
- Phil Wang's vector-quantize-pytorch implementation in PyTorch.
- Origianl source: https://github.com/lucidrains/vector-quantize-pytorch
+ Original source: https://github.com/lucidrains/vector-quantize-pytorch
- The vector-quantize-pytorch code is licensed under the MIT License:
@@ -419,9 +419,9 @@ class ResidualVQ(nn.Module):
# and the network should be able to reconstruct
if quantize_dim < self.num_quantizers:
- assert (
- self.quantize_dropout > 0.0
- ), "quantize dropout must be greater than 0 if you wish to reconstruct from a signal with less fine quantizations"
+ assert self.quantize_dropout > 0.0, (
+ "quantize dropout must be greater than 0 if you wish to reconstruct from a signal with less fine quantizations"
+ )
indices = F.pad(indices, (0, self.num_quantizers - quantize_dim), value=-1)
# get ready for gathering
@@ -472,9 +472,9 @@ class ResidualVQ(nn.Module):
all_indices = []
if return_loss:
- assert not torch.any(
- indices == -1
- ), "some of the residual vq indices were dropped out. please use indices derived when the module is in eval mode to derive cross entropy loss"
+ assert not torch.any(indices == -1), (
+ "some of the residual vq indices were dropped out. please use indices derived when the module is in eval mode to derive cross entropy loss"
+ )
ce_losses = []
should_quantize_dropout = self.training and self.quantize_dropout and not return_loss
@@ -887,9 +887,9 @@ class VectorQuantize(nn.Module):
# only calculate orthogonal loss for the activated codes for this batch
if self.orthogonal_reg_active_codes_only:
- assert not (
- is_multiheaded and self.separate_codebook_per_head
- ), "orthogonal regularization for only active codes not compatible with multi-headed with separate codebooks yet"
+ assert not (is_multiheaded and self.separate_codebook_per_head), (
+ "orthogonal regularization for only active codes not compatible with multi-headed with separate codebooks yet"
+ )
unique_code_ids = torch.unique(embed_ind)
codebook = codebook[:, unique_code_ids]
@@ -999,9 +999,9 @@ def gumbel_sample(
ind = sampling_logits.argmax(dim=dim)
one_hot = F.one_hot(ind, size).type(dtype)
- assert not (
- reinmax and not straight_through
- ), "reinmax can only be turned on if using straight through gumbel softmax"
+ assert not (reinmax and not straight_through), (
+ "reinmax can only be turned on if using straight through gumbel softmax"
+ )
if not straight_through or temperature <= 0.0 or not training:
return ind, one_hot
@@ -1209,9 +1209,9 @@ class EuclideanCodebook(nn.Module):
self.gumbel_sample = gumbel_sample
self.sample_codebook_temp = sample_codebook_temp
- assert not (
- use_ddp and num_codebooks > 1 and kmeans_init
- ), "kmeans init is not compatible with multiple codebooks in distributed environment for now"
+ assert not (use_ddp and num_codebooks > 1 and kmeans_init), (
+ "kmeans init is not compatible with multiple codebooks in distributed environment for now"
+ )
self.sample_fn = sample_vectors_distributed if use_ddp and sync_kmeans else batched_sample_vectors
self.kmeans_all_reduce_fn = distributed.all_reduce if use_ddp and sync_kmeans else noop
@@ -1349,9 +1349,9 @@ class EuclideanCodebook(nn.Module):
# calculate distributed variance
- variance_numer = reduce((data - batch_mean) ** 2, "h n d -> h 1 d", "sum")
- distributed.all_reduce(variance_numer)
- batch_variance = variance_numer / num_vectors
+ variance_number = reduce((data - batch_mean) ** 2, "h n d -> h 1 d", "sum")
+ distributed.all_reduce(variance_number)
+ batch_variance = variance_number / num_vectors
self.update_with_decay("batch_variance", batch_variance, self.affine_param_batch_decay)
diff --git a/lerobot/common/robot_devices/cameras/configs.py b/lerobot/common/robot_devices/cameras/configs.py
index 6acdbd3e..013419a9 100644
--- a/lerobot/common/robot_devices/cameras/configs.py
+++ b/lerobot/common/robot_devices/cameras/configs.py
@@ -1,3 +1,17 @@
+# 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 abc
from dataclasses import dataclass
diff --git a/lerobot/common/robot_devices/cameras/intelrealsense.py b/lerobot/common/robot_devices/cameras/intelrealsense.py
index 7e65dba9..7a21661a 100644
--- a/lerobot/common/robot_devices/cameras/intelrealsense.py
+++ b/lerobot/common/robot_devices/cameras/intelrealsense.py
@@ -1,3 +1,17 @@
+# 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.
+
"""
This file contains utilities for recording frames from Intel Realsense cameras.
"""
@@ -34,7 +48,7 @@ def find_cameras(raise_when_empty=True, mock=False) -> list[dict]:
connected to the computer.
"""
if mock:
- import tests.mock_pyrealsense2 as rs
+ import tests.cameras.mock_pyrealsense2 as rs
else:
import pyrealsense2 as rs
@@ -86,7 +100,7 @@ def save_images_from_cameras(
serial_numbers = [cam["serial_number"] for cam in camera_infos]
if mock:
- import tests.mock_cv2 as cv2
+ import tests.cameras.mock_cv2 as cv2
else:
import cv2
@@ -100,7 +114,7 @@ def save_images_from_cameras(
camera = IntelRealSenseCamera(config)
camera.connect()
print(
- f"IntelRealSenseCamera({camera.serial_number}, fps={camera.fps}, width={camera.width}, height={camera.height}, color_mode={camera.color_mode})"
+ f"IntelRealSenseCamera({camera.serial_number}, fps={camera.fps}, width={camera.capture_width}, height={camera.capture_height}, color_mode={camera.color_mode})"
)
cameras.append(camera)
@@ -210,9 +224,20 @@ class IntelRealSenseCamera:
self.serial_number = self.find_serial_number_from_name(config.name)
else:
self.serial_number = config.serial_number
+
+ # Store the raw (capture) resolution from the config.
+ self.capture_width = config.width
+ self.capture_height = config.height
+
+ # If rotated by ±90, swap width and height.
+ if config.rotation in [-90, 90]:
+ self.width = config.height
+ self.height = config.width
+ else:
+ self.width = config.width
+ self.height = config.height
+
self.fps = config.fps
- self.width = config.width
- self.height = config.height
self.channels = config.channels
self.color_mode = config.color_mode
self.use_depth = config.use_depth
@@ -228,11 +253,10 @@ class IntelRealSenseCamera:
self.logs = {}
if self.mock:
- import tests.mock_cv2 as cv2
+ import tests.cameras.mock_cv2 as cv2
else:
import cv2
- # TODO(alibets): Do we keep original width/height or do we define them after rotation?
self.rotation = None
if config.rotation == -90:
self.rotation = cv2.ROTATE_90_COUNTERCLOCKWISE
@@ -263,22 +287,26 @@ class IntelRealSenseCamera:
)
if self.mock:
- import tests.mock_pyrealsense2 as rs
+ import tests.cameras.mock_pyrealsense2 as rs
else:
import pyrealsense2 as rs
config = rs.config()
config.enable_device(str(self.serial_number))
- if self.fps and self.width and self.height:
+ if self.fps and self.capture_width and self.capture_height:
# TODO(rcadene): can we set rgb8 directly?
- config.enable_stream(rs.stream.color, self.width, self.height, rs.format.rgb8, self.fps)
+ config.enable_stream(
+ rs.stream.color, self.capture_width, self.capture_height, rs.format.rgb8, self.fps
+ )
else:
config.enable_stream(rs.stream.color)
if self.use_depth:
- if self.fps and self.width and self.height:
- config.enable_stream(rs.stream.depth, self.width, self.height, rs.format.z16, self.fps)
+ if self.fps and self.capture_width and self.capture_height:
+ config.enable_stream(
+ rs.stream.depth, self.capture_width, self.capture_height, rs.format.z16, self.fps
+ )
else:
config.enable_stream(rs.stream.depth)
@@ -316,18 +344,18 @@ class IntelRealSenseCamera:
raise OSError(
f"Can't set {self.fps=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_fps}."
)
- if self.width is not None and self.width != actual_width:
+ if self.capture_width is not None and self.capture_width != actual_width:
raise OSError(
- f"Can't set {self.width=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_width}."
+ f"Can't set {self.capture_width=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_width}."
)
- if self.height is not None and self.height != actual_height:
+ if self.capture_height is not None and self.capture_height != actual_height:
raise OSError(
- f"Can't set {self.height=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_height}."
+ f"Can't set {self.capture_height=} for IntelRealSenseCamera({self.serial_number}). Actual value is {actual_height}."
)
self.fps = round(actual_fps)
- self.width = round(actual_width)
- self.height = round(actual_height)
+ self.capture_width = round(actual_width)
+ self.capture_height = round(actual_height)
self.is_connected = True
@@ -347,7 +375,7 @@ class IntelRealSenseCamera:
)
if self.mock:
- import tests.mock_cv2 as cv2
+ import tests.cameras.mock_cv2 as cv2
else:
import cv2
@@ -373,7 +401,7 @@ class IntelRealSenseCamera:
color_image = cv2.cvtColor(color_image, cv2.COLOR_RGB2BGR)
h, w, _ = color_image.shape
- if h != self.height or w != self.width:
+ if h != self.capture_height or w != self.capture_width:
raise OSError(
f"Can't capture color image with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
)
@@ -395,7 +423,7 @@ class IntelRealSenseCamera:
depth_map = np.asanyarray(depth_frame.get_data())
h, w = depth_map.shape
- if h != self.height or w != self.width:
+ if h != self.capture_height or w != self.capture_width:
raise OSError(
f"Can't capture depth map with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
)
diff --git a/lerobot/common/robot_devices/cameras/opencv.py b/lerobot/common/robot_devices/cameras/opencv.py
index 93c791fa..f279f315 100644
--- a/lerobot/common/robot_devices/cameras/opencv.py
+++ b/lerobot/common/robot_devices/cameras/opencv.py
@@ -1,3 +1,17 @@
+# 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.
+
"""
This file contains utilities for recording frames from cameras. For more info look at `OpenCVCamera` docstring.
"""
@@ -66,7 +80,7 @@ def _find_cameras(
possible_camera_ids: list[int | str], raise_when_empty=False, mock=False
) -> list[int | str]:
if mock:
- import tests.mock_cv2 as cv2
+ import tests.cameras.mock_cv2 as cv2
else:
import cv2
@@ -130,8 +144,8 @@ def save_images_from_cameras(
camera = OpenCVCamera(config)
camera.connect()
print(
- f"OpenCVCamera({camera.camera_index}, fps={camera.fps}, width={camera.width}, "
- f"height={camera.height}, color_mode={camera.color_mode})"
+ f"OpenCVCamera({camera.camera_index}, fps={camera.fps}, width={camera.capture_width}, "
+ f"height={camera.capture_height}, color_mode={camera.color_mode})"
)
cameras.append(camera)
@@ -230,9 +244,19 @@ class OpenCVCamera:
else:
raise ValueError(f"Please check the provided camera_index: {self.camera_index}")
+ # Store the raw (capture) resolution from the config.
+ self.capture_width = config.width
+ self.capture_height = config.height
+
+ # If rotated by ±90, swap width and height.
+ if config.rotation in [-90, 90]:
+ self.width = config.height
+ self.height = config.width
+ else:
+ self.width = config.width
+ self.height = config.height
+
self.fps = config.fps
- self.width = config.width
- self.height = config.height
self.channels = config.channels
self.color_mode = config.color_mode
self.mock = config.mock
@@ -245,11 +269,10 @@ class OpenCVCamera:
self.logs = {}
if self.mock:
- import tests.mock_cv2 as cv2
+ import tests.cameras.mock_cv2 as cv2
else:
import cv2
- # TODO(aliberts): Do we keep original width/height or do we define them after rotation?
self.rotation = None
if config.rotation == -90:
self.rotation = cv2.ROTATE_90_COUNTERCLOCKWISE
@@ -263,7 +286,7 @@ class OpenCVCamera:
raise RobotDeviceAlreadyConnectedError(f"OpenCVCamera({self.camera_index}) is already connected.")
if self.mock:
- import tests.mock_cv2 as cv2
+ import tests.cameras.mock_cv2 as cv2
else:
import cv2
@@ -271,10 +294,20 @@ class OpenCVCamera:
# when other threads are used to save the images.
cv2.setNumThreads(1)
+ backend = (
+ cv2.CAP_V4L2
+ if platform.system() == "Linux"
+ else cv2.CAP_DSHOW
+ if platform.system() == "Windows"
+ else cv2.CAP_AVFOUNDATION
+ if platform.system() == "Darwin"
+ else cv2.CAP_ANY
+ )
+
camera_idx = f"/dev/video{self.camera_index}" if platform.system() == "Linux" else self.camera_index
# First create a temporary camera trying to access `camera_index`,
# and verify it is a valid camera by calling `isOpened`.
- tmp_camera = cv2.VideoCapture(camera_idx)
+ tmp_camera = cv2.VideoCapture(camera_idx, backend)
is_camera_open = tmp_camera.isOpened()
# Release camera to make it accessible for `find_camera_indices`
tmp_camera.release()
@@ -297,14 +330,14 @@ class OpenCVCamera:
# Secondly, create the camera that will be used downstream.
# Note: For some unknown reason, calling `isOpened` blocks the camera which then
# needs to be re-created.
- self.camera = cv2.VideoCapture(camera_idx)
+ self.camera = cv2.VideoCapture(camera_idx, backend)
if self.fps is not None:
self.camera.set(cv2.CAP_PROP_FPS, self.fps)
- if self.width is not None:
- self.camera.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
- if self.height is not None:
- self.camera.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
+ if self.capture_width is not None:
+ self.camera.set(cv2.CAP_PROP_FRAME_WIDTH, self.capture_width)
+ if self.capture_height is not None:
+ self.camera.set(cv2.CAP_PROP_FRAME_HEIGHT, self.capture_height)
actual_fps = self.camera.get(cv2.CAP_PROP_FPS)
actual_width = self.camera.get(cv2.CAP_PROP_FRAME_WIDTH)
@@ -316,19 +349,22 @@ class OpenCVCamera:
raise OSError(
f"Can't set {self.fps=} for OpenCVCamera({self.camera_index}). Actual value is {actual_fps}."
)
- if self.width is not None and not math.isclose(self.width, actual_width, rel_tol=1e-3):
+ if self.capture_width is not None and not math.isclose(
+ self.capture_width, actual_width, rel_tol=1e-3
+ ):
raise OSError(
- f"Can't set {self.width=} for OpenCVCamera({self.camera_index}). Actual value is {actual_width}."
+ f"Can't set {self.capture_width=} for OpenCVCamera({self.camera_index}). Actual value is {actual_width}."
)
- if self.height is not None and not math.isclose(self.height, actual_height, rel_tol=1e-3):
+ if self.capture_height is not None and not math.isclose(
+ self.capture_height, actual_height, rel_tol=1e-3
+ ):
raise OSError(
- f"Can't set {self.height=} for OpenCVCamera({self.camera_index}). Actual value is {actual_height}."
+ f"Can't set {self.capture_height=} for OpenCVCamera({self.camera_index}). Actual value is {actual_height}."
)
self.fps = round(actual_fps)
- self.width = round(actual_width)
- self.height = round(actual_height)
-
+ self.capture_width = round(actual_width)
+ self.capture_height = round(actual_height)
self.is_connected = True
def read(self, temporary_color_mode: str | None = None) -> np.ndarray:
@@ -362,14 +398,14 @@ class OpenCVCamera:
# so we convert the image color from BGR to RGB.
if requested_color_mode == "rgb":
if self.mock:
- import tests.mock_cv2 as cv2
+ import tests.cameras.mock_cv2 as cv2
else:
import cv2
color_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2RGB)
h, w, _ = color_image.shape
- if h != self.height or w != self.width:
+ if h != self.capture_height or w != self.capture_width:
raise OSError(
f"Can't capture color image with expected height and width ({self.height} x {self.width}). ({h} x {w}) returned instead."
)
diff --git a/lerobot/common/robot_devices/cameras/utils.py b/lerobot/common/robot_devices/cameras/utils.py
index 88288ea3..c6431646 100644
--- a/lerobot/common/robot_devices/cameras/utils.py
+++ b/lerobot/common/robot_devices/cameras/utils.py
@@ -1,3 +1,17 @@
+# 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.
+
from typing import Protocol
import numpy as np
@@ -31,7 +45,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
diff --git a/lerobot/common/robot_devices/control_configs.py b/lerobot/common/robot_devices/control_configs.py
index a2f3889c..0ecd8683 100644
--- a/lerobot/common/robot_devices/control_configs.py
+++ b/lerobot/common/robot_devices/control_configs.py
@@ -1,14 +1,25 @@
-import logging
+# 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.
+
from dataclasses import dataclass
from pathlib import Path
import draccus
from lerobot.common.robot_devices.robots.configs import RobotConfig
-from lerobot.common.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
from lerobot.configs import parser
from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.train import TrainPipelineConfig
@dataclass
@@ -43,11 +54,6 @@ class RecordControlConfig(ControlConfig):
# Root directory where the dataset will be stored (e.g. 'dataset/path').
root: str | Path | None = None
policy: PreTrainedConfig | None = None
- # TODO(rcadene, aliberts): By default, use device and use_amp values from policy checkpoint.
- device: str | None = None # cuda | cpu | mps
- # `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
- # automatic gradient scaling is used.
- use_amp: bool | None = None
# Limit the frames per second. By default, uses the policy fps.
fps: int | None = None
# Number of seconds before starting data collection. It allows the robot devices to warmup and synchronize.
@@ -60,15 +66,13 @@ class RecordControlConfig(ControlConfig):
num_episodes: int = 50
# Encode frames in the dataset into video
video: bool = True
- # By default, run the computation of the data statistics at the end of data collection. Compute intensive and not required to just replay an episode.
- run_compute_stats: bool = True
# Upload dataset to Hugging Face hub.
push_to_hub: bool = True
# Upload on private repository on the Hugging Face hub.
private: bool = False
# Add tags to your dataset on the hub.
tags: list[str] | None = None
- # Number of subprocesses handling the saving of frames as PNGs. Set to 0 to use threads only;
+ # Number of subprocesses handling the saving of frames as PNG. Set to 0 to use threads only;
# set to ≥1 to use subprocesses, each using threads to write images. The best number of processes
# and threads depends on your system. We recommend 4 threads per camera with 0 processes.
# If fps is unstable, adjust the thread count. If still unstable, try using 1 or more subprocesses.
@@ -83,9 +87,6 @@ class RecordControlConfig(ControlConfig):
play_sounds: bool = True
# Resume recording on an existing dataset.
resume: bool = False
- # TODO(rcadene, aliberts): remove local_files_only when refactor with dataset as argument
- # Use local files only. By default, this script will try to fetch the dataset from the hub if it exists.
- local_files_only: bool = False
def __post_init__(self):
# HACK: We parse again the cli args here to get the pretrained path if there was one.
@@ -95,27 +96,6 @@ class RecordControlConfig(ControlConfig):
self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
self.policy.pretrained_path = policy_path
- # When no device or use_amp are given, use the one from training config.
- if self.device is None or self.use_amp is None:
- train_cfg = TrainPipelineConfig.from_pretrained(policy_path)
- if self.device is None:
- self.device = train_cfg.device
- if self.use_amp is None:
- self.use_amp = train_cfg.use_amp
-
- # Automatically switch to available device if necessary
- if not is_torch_device_available(self.device):
- auto_device = auto_select_torch_device()
- logging.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
- self.device = auto_device
-
- # Automatically deactivate AMP if necessary
- if self.use_amp and not is_amp_available(self.device):
- logging.warning(
- f"Automatic Mixed Precision (amp) is not available on device '{self.device}'. Deactivating AMP."
- )
- self.use_amp = False
-
@ControlConfig.register_subclass("replay")
@dataclass
@@ -130,9 +110,12 @@ class ReplayControlConfig(ControlConfig):
fps: int | None = None
# Use vocal synthesis to read events.
play_sounds: bool = True
- # TODO(rcadene, aliberts): remove local_files_only when refactor with dataset as argument
- # Use local files only. By default, this script will try to fetch the dataset from the hub if it exists.
- local_files_only: bool = False
+
+
+@ControlConfig.register_subclass("remote_robot")
+@dataclass
+class RemoteRobotConfig(ControlConfig):
+ log_interval: int = 100
@dataclass
diff --git a/lerobot/common/robot_devices/control_utils.py b/lerobot/common/robot_devices/control_utils.py
index 9368b89d..78a8c6a6 100644
--- a/lerobot/common/robot_devices/control_utils.py
+++ b/lerobot/common/robot_devices/control_utils.py
@@ -1,3 +1,17 @@
+# 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.
+
########################################################################################
# Utilities
########################################################################################
@@ -12,13 +26,13 @@ from functools import cache
import cv2
import torch
-import tqdm
from deepdiff import DeepDiff
from termcolor import colored
from lerobot.common.datasets.image_writer import safe_stop_image_writer
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
from lerobot.common.datasets.utils import get_features_from_robot
+from lerobot.common.policies.pretrained import PreTrainedPolicy
from lerobot.common.robot_devices.robots.utils import Robot
from lerobot.common.robot_devices.utils import busy_wait
from lerobot.common.utils.utils import get_safe_torch_device, has_method
@@ -33,7 +47,7 @@ def log_control_info(robot: Robot, dt_s, episode_index=None, frame_index=None, f
def log_dt(shortname, dt_val_s):
nonlocal log_items, fps
- info_str = f"{shortname}:{dt_val_s * 1000:5.2f} ({1/ dt_val_s:3.1f}hz)"
+ info_str = f"{shortname}:{dt_val_s * 1000:5.2f} ({1 / dt_val_s:3.1f}hz)"
if fps is not None:
actual_fps = 1 / dt_val_s
if actual_fps < fps - 1:
@@ -180,9 +194,8 @@ def record_episode(
episode_time_s,
display_cameras,
policy,
- device,
- use_amp,
fps,
+ single_task,
):
control_loop(
robot=robot,
@@ -191,10 +204,9 @@ def record_episode(
dataset=dataset,
events=events,
policy=policy,
- device=device,
- use_amp=use_amp,
fps=fps,
teleoperate=policy is None,
+ single_task=single_task,
)
@@ -206,10 +218,9 @@ def control_loop(
display_cameras=False,
dataset: LeRobotDataset | None = None,
events=None,
- policy=None,
- device: torch.device | str | None = None,
- use_amp: bool | None = None,
+ policy: PreTrainedPolicy = None,
fps: int | None = None,
+ single_task: str | None = None,
):
# TODO(rcadene): Add option to record logs
if not robot.is_connected:
@@ -224,12 +235,12 @@ def control_loop(
if teleoperate and policy is not None:
raise ValueError("When `teleoperate` is True, `policy` should be None.")
+ if dataset is not None and single_task is None:
+ raise ValueError("You need to provide a task as argument in `single_task`.")
+
if dataset is not None and fps is not None and dataset.fps != fps:
raise ValueError(f"The dataset fps should be equal to requested fps ({dataset['fps']} != {fps}).")
- if isinstance(device, str):
- device = get_safe_torch_device(device)
-
timestamp = 0
start_episode_t = time.perf_counter()
while timestamp < control_time_s:
@@ -241,14 +252,16 @@ def control_loop(
observation = robot.capture_observation()
if policy is not None:
- pred_action = predict_action(observation, policy, device, use_amp)
+ pred_action = predict_action(
+ observation, policy, get_safe_torch_device(policy.config.device), policy.config.use_amp
+ )
# Action can eventually be clipped using `max_relative_target`,
# so action actually sent is saved in the dataset.
action = robot.send_action(pred_action)
action = {"action": action}
if dataset is not None:
- frame = {**observation, **action}
+ frame = {**observation, **action, "task": single_task}
dataset.add_frame(frame)
if display_cameras and not is_headless():
@@ -270,24 +283,18 @@ def control_loop(
break
-def reset_environment(robot, events, reset_time_s):
+def reset_environment(robot, events, reset_time_s, fps):
# TODO(rcadene): refactor warmup_record and reset_environment
- # TODO(alibets): allow for teleop during reset
if has_method(robot, "teleop_safety_stop"):
robot.teleop_safety_stop()
- timestamp = 0
- start_vencod_t = time.perf_counter()
-
- # Wait if necessary
- with tqdm.tqdm(total=reset_time_s, desc="Waiting") as pbar:
- while timestamp < reset_time_s:
- time.sleep(1)
- timestamp = time.perf_counter() - start_vencod_t
- pbar.update(1)
- if events["exit_early"]:
- events["exit_early"] = False
- break
+ control_loop(
+ robot=robot,
+ control_time_s=reset_time_s,
+ events=events,
+ fps=fps,
+ teleoperate=True,
+ )
def stop_recording(robot, listener, display_cameras):
diff --git a/lerobot/common/robot_devices/motors/configs.py b/lerobot/common/robot_devices/motors/configs.py
index 37b781f9..0bfbaf83 100644
--- a/lerobot/common/robot_devices/motors/configs.py
+++ b/lerobot/common/robot_devices/motors/configs.py
@@ -1,3 +1,17 @@
+# 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 abc
from dataclasses import dataclass
diff --git a/lerobot/common/robot_devices/motors/dynamixel.py b/lerobot/common/robot_devices/motors/dynamixel.py
index 54836d8e..6096ceb5 100644
--- a/lerobot/common/robot_devices/motors/dynamixel.py
+++ b/lerobot/common/robot_devices/motors/dynamixel.py
@@ -1,3 +1,17 @@
+# 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 enum
import logging
import math
@@ -242,7 +256,7 @@ class DriveMode(enum.Enum):
class CalibrationMode(enum.Enum):
# Joints with rotational motions are expressed in degrees in nominal range of [-180, 180]
DEGREE = 0
- # Joints with linear motions (like gripper of Aloha) are experessed in nominal range of [0, 100]
+ # Joints with linear motions (like gripper of Aloha) are expressed in nominal range of [0, 100]
LINEAR = 1
@@ -318,7 +332,7 @@ class DynamixelMotorsBus:
)
if self.mock:
- import tests.mock_dynamixel_sdk as dxl
+ import tests.motors.mock_dynamixel_sdk as dxl
else:
import dynamixel_sdk as dxl
@@ -342,7 +356,7 @@ class DynamixelMotorsBus:
def reconnect(self):
if self.mock:
- import tests.mock_dynamixel_sdk as dxl
+ import tests.motors.mock_dynamixel_sdk as dxl
else:
import dynamixel_sdk as dxl
@@ -610,7 +624,7 @@ class DynamixelMotorsBus:
# 0-centered resolution range (e.g. [-2048, 2048] for resolution=4096)
values[i] = values[i] / HALF_TURN_DEGREE * (resolution // 2)
- # Substract the homing offsets to come back to actual motor range of values
+ # Subtract the homing offsets to come back to actual motor range of values
# which can be arbitrary.
values[i] -= homing_offset
@@ -632,7 +646,7 @@ class DynamixelMotorsBus:
def read_with_motor_ids(self, motor_models, motor_ids, data_name, num_retry=NUM_READ_RETRY):
if self.mock:
- import tests.mock_dynamixel_sdk as dxl
+ import tests.motors.mock_dynamixel_sdk as dxl
else:
import dynamixel_sdk as dxl
@@ -677,7 +691,7 @@ class DynamixelMotorsBus:
start_time = time.perf_counter()
if self.mock:
- import tests.mock_dynamixel_sdk as dxl
+ import tests.motors.mock_dynamixel_sdk as dxl
else:
import dynamixel_sdk as dxl
@@ -743,7 +757,7 @@ class DynamixelMotorsBus:
def write_with_motor_ids(self, motor_models, motor_ids, data_name, values, num_retry=NUM_WRITE_RETRY):
if self.mock:
- import tests.mock_dynamixel_sdk as dxl
+ import tests.motors.mock_dynamixel_sdk as dxl
else:
import dynamixel_sdk as dxl
@@ -779,7 +793,7 @@ class DynamixelMotorsBus:
start_time = time.perf_counter()
if self.mock:
- import tests.mock_dynamixel_sdk as dxl
+ import tests.motors.mock_dynamixel_sdk as dxl
else:
import dynamixel_sdk as dxl
diff --git a/lerobot/common/robot_devices/motors/feetech.py b/lerobot/common/robot_devices/motors/feetech.py
index a59db7df..64c7f413 100644
--- a/lerobot/common/robot_devices/motors/feetech.py
+++ b/lerobot/common/robot_devices/motors/feetech.py
@@ -1,3 +1,17 @@
+# 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 enum
import logging
import math
@@ -221,7 +235,7 @@ class DriveMode(enum.Enum):
class CalibrationMode(enum.Enum):
# Joints with rotational motions are expressed in degrees in nominal range of [-180, 180]
DEGREE = 0
- # Joints with linear motions (like gripper of Aloha) are experessed in nominal range of [0, 100]
+ # Joints with linear motions (like gripper of Aloha) are expressed in nominal range of [0, 100]
LINEAR = 1
@@ -299,7 +313,7 @@ class FeetechMotorsBus:
)
if self.mock:
- import tests.mock_scservo_sdk as scs
+ import tests.motors.mock_scservo_sdk as scs
else:
import scservo_sdk as scs
@@ -323,7 +337,7 @@ class FeetechMotorsBus:
def reconnect(self):
if self.mock:
- import tests.mock_scservo_sdk as scs
+ import tests.motors.mock_scservo_sdk as scs
else:
import scservo_sdk as scs
@@ -591,7 +605,7 @@ class FeetechMotorsBus:
# 0-centered resolution range (e.g. [-2048, 2048] for resolution=4096)
values[i] = values[i] / HALF_TURN_DEGREE * (resolution // 2)
- # Substract the homing offsets to come back to actual motor range of values
+ # Subtract the homing offsets to come back to actual motor range of values
# which can be arbitrary.
values[i] -= homing_offset
@@ -632,7 +646,7 @@ class FeetechMotorsBus:
track["prev"][idx] = values[i]
continue
- # Detect a full rotation occured
+ # Detect a full rotation occurred
if abs(track["prev"][idx] - values[i]) > 2048:
# Position went below 0 and got reset to 4095
if track["prev"][idx] < values[i]:
@@ -650,7 +664,7 @@ class FeetechMotorsBus:
def read_with_motor_ids(self, motor_models, motor_ids, data_name, num_retry=NUM_READ_RETRY):
if self.mock:
- import tests.mock_scservo_sdk as scs
+ import tests.motors.mock_scservo_sdk as scs
else:
import scservo_sdk as scs
@@ -688,7 +702,7 @@ class FeetechMotorsBus:
def read(self, data_name, motor_names: str | list[str] | None = None):
if self.mock:
- import tests.mock_scservo_sdk as scs
+ import tests.motors.mock_scservo_sdk as scs
else:
import scservo_sdk as scs
@@ -768,7 +782,7 @@ class FeetechMotorsBus:
def write_with_motor_ids(self, motor_models, motor_ids, data_name, values, num_retry=NUM_WRITE_RETRY):
if self.mock:
- import tests.mock_scservo_sdk as scs
+ import tests.motors.mock_scservo_sdk as scs
else:
import scservo_sdk as scs
@@ -804,7 +818,7 @@ class FeetechMotorsBus:
start_time = time.perf_counter()
if self.mock:
- import tests.mock_scservo_sdk as scs
+ import tests.motors.mock_scservo_sdk as scs
else:
import scservo_sdk as scs
diff --git a/lerobot/common/robot_devices/motors/utils.py b/lerobot/common/robot_devices/motors/utils.py
index fc64f050..bd86f4c6 100644
--- a/lerobot/common/robot_devices/motors/utils.py
+++ b/lerobot/common/robot_devices/motors/utils.py
@@ -1,3 +1,17 @@
+# 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.
+
from typing import Protocol
from lerobot.common.robot_devices.motors.configs import (
diff --git a/lerobot/common/robot_devices/robots/configs.py b/lerobot/common/robot_devices/robots/configs.py
index a976f601..e940b442 100644
--- a/lerobot/common/robot_devices/robots/configs.py
+++ b/lerobot/common/robot_devices/robots/configs.py
@@ -1,3 +1,17 @@
+# 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 abc
from dataclasses import dataclass, field
from typing import Sequence
@@ -514,3 +528,86 @@ class StretchRobotConfig(RobotConfig):
)
mock: bool = False
+
+
+@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 = "192.168.0.193"
+ port: int = 5555
+ video_port: int = 5556
+
+ cameras: dict[str, CameraConfig] = field(
+ default_factory=lambda: {
+ "front": OpenCVCameraConfig(
+ camera_index="/dev/video0", fps=30, width=640, height=480, rotation=90
+ ),
+ "wrist": OpenCVCameraConfig(
+ camera_index="/dev/video2", 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/ttyACM0",
+ 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
diff --git a/lerobot/common/robot_devices/robots/dynamixel_calibration.py b/lerobot/common/robot_devices/robots/dynamixel_calibration.py
index 5c4932d2..98fe8754 100644
--- a/lerobot/common/robot_devices/robots/dynamixel_calibration.py
+++ b/lerobot/common/robot_devices/robots/dynamixel_calibration.py
@@ -1,3 +1,17 @@
+# 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.
+
"""Logic to calibrate a robot arm built with dynamixel motors"""
# TODO(rcadene, aliberts): move this logic into the robot code when refactoring
@@ -87,7 +101,7 @@ def run_arm_calibration(arm: MotorsBus, robot_type: str, arm_name: str, arm_type
# For instance, if the motor rotates 90 degree, and its value is -90 after applying the homing offset, then we know its rotation direction
# is inverted. However, for the calibration being successful, we need everyone to follow the same target position.
# Sometimes, there is only one possible rotation direction. For instance, if the gripper is closed, there is only one direction which
- # corresponds to opening the gripper. When the rotation direction is ambiguous, we arbitrarely rotate clockwise from the point of view
+ # corresponds to opening the gripper. When the rotation direction is ambiguous, we arbitrarily rotate clockwise from the point of view
# of the previous motor in the kinetic chain.
print("\nMove arm to rotated target position")
print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="rotated"))
@@ -115,7 +129,7 @@ def run_arm_calibration(arm: MotorsBus, robot_type: str, arm_name: str, arm_type
# TODO(rcadene): make type of joints (DEGREE or LINEAR) configurable from yaml?
if robot_type in ["aloha"] and "gripper" in arm.motor_names:
- # Joints with linear motions (like gripper of Aloha) are experessed in nominal range of [0, 100]
+ # Joints with linear motions (like gripper of Aloha) are expressed in nominal range of [0, 100]
calib_idx = arm.motor_names.index("gripper")
calib_mode[calib_idx] = CalibrationMode.LINEAR.name
diff --git a/lerobot/common/robot_devices/robots/feetech_calibration.py b/lerobot/common/robot_devices/robots/feetech_calibration.py
index b015951a..2c1e7180 100644
--- a/lerobot/common/robot_devices/robots/feetech_calibration.py
+++ b/lerobot/common/robot_devices/robots/feetech_calibration.py
@@ -1,3 +1,17 @@
+# 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.
+
"""Logic to calibrate a robot arm built with feetech motors"""
# TODO(rcadene, aliberts): move this logic into the robot code when refactoring
@@ -443,7 +457,7 @@ def run_arm_manual_calibration(arm: MotorsBus, robot_type: str, arm_name: str, a
# For instance, if the motor rotates 90 degree, and its value is -90 after applying the homing offset, then we know its rotation direction
# is inverted. However, for the calibration being successful, we need everyone to follow the same target position.
# Sometimes, there is only one possible rotation direction. For instance, if the gripper is closed, there is only one direction which
- # corresponds to opening the gripper. When the rotation direction is ambiguous, we arbitrarely rotate clockwise from the point of view
+ # corresponds to opening the gripper. When the rotation direction is ambiguous, we arbitrarily rotate clockwise from the point of view
# of the previous motor in the kinetic chain.
print("\nMove arm to rotated target position")
print("See: " + URL_TEMPLATE.format(robot=robot_type, arm=arm_type, position="rotated"))
diff --git a/lerobot/common/robot_devices/robots/lekiwi_remote.py b/lerobot/common/robot_devices/robots/lekiwi_remote.py
new file mode 100644
index 00000000..7bf52d21
--- /dev/null
+++ b/lerobot/common/robot_devices/robots/lekiwi_remote.py
@@ -0,0 +1,224 @@
+# 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 base64
+import json
+import threading
+import time
+from pathlib import Path
+
+import cv2
+import zmq
+
+from lerobot.common.robot_devices.robots.mobile_manipulator import LeKiwi
+
+
+def setup_zmq_sockets(config):
+ context = zmq.Context()
+ cmd_socket = context.socket(zmq.PULL)
+ cmd_socket.setsockopt(zmq.CONFLATE, 1)
+ cmd_socket.bind(f"tcp://*:{config.port}")
+
+ video_socket = context.socket(zmq.PUSH)
+ video_socket.setsockopt(zmq.CONFLATE, 1)
+ video_socket.bind(f"tcp://*:{config.video_port}")
+
+ return context, cmd_socket, video_socket
+
+
+def run_camera_capture(cameras, images_lock, latest_images_dict, stop_event):
+ while not stop_event.is_set():
+ local_dict = {}
+ for name, cam in cameras.items():
+ frame = cam.async_read()
+ ret, buffer = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
+ if ret:
+ local_dict[name] = base64.b64encode(buffer).decode("utf-8")
+ else:
+ local_dict[name] = ""
+ with images_lock:
+ latest_images_dict.update(local_dict)
+ time.sleep(0.01)
+
+
+def calibrate_follower_arm(motors_bus, calib_dir_str):
+ """
+ Calibrates the follower arm. Attempts to load an existing calibration file;
+ if not found, runs manual calibration and saves the result.
+ """
+ calib_dir = Path(calib_dir_str)
+ calib_dir.mkdir(parents=True, exist_ok=True)
+ calib_file = calib_dir / "main_follower.json"
+ try:
+ from lerobot.common.robot_devices.robots.feetech_calibration import run_arm_manual_calibration
+ except ImportError:
+ print("[WARNING] Calibration function not available. Skipping calibration.")
+ return
+
+ if calib_file.exists():
+ with open(calib_file) as f:
+ calibration = json.load(f)
+ print(f"[INFO] Loaded calibration from {calib_file}")
+ else:
+ print("[INFO] Calibration file not found. Running manual calibration...")
+ calibration = run_arm_manual_calibration(motors_bus, "lekiwi", "follower_arm", "follower")
+ print(f"[INFO] Calibration complete. Saving to {calib_file}")
+ with open(calib_file, "w") as f:
+ json.dump(calibration, f)
+ try:
+ motors_bus.set_calibration(calibration)
+ print("[INFO] Applied calibration for follower arm.")
+ except Exception as e:
+ print(f"[WARNING] Could not apply calibration: {e}")
+
+
+def run_lekiwi(robot_config):
+ """
+ Runs the LeKiwi robot:
+ - Sets up cameras and connects them.
+ - Initializes the follower arm motors.
+ - Calibrates the follower arm if necessary.
+ - Creates ZeroMQ sockets for receiving commands and streaming observations.
+ - Processes incoming commands (arm and wheel commands) and sends back sensor and camera data.
+ """
+ # Import helper functions and classes
+ from lerobot.common.robot_devices.cameras.utils import make_cameras_from_configs
+ from lerobot.common.robot_devices.motors.feetech import FeetechMotorsBus, TorqueMode
+
+ # Initialize cameras from the robot configuration.
+ cameras = make_cameras_from_configs(robot_config.cameras)
+ for cam in cameras.values():
+ cam.connect()
+
+ # Initialize the motors bus using the follower arm configuration.
+ motor_config = robot_config.follower_arms.get("main")
+ if motor_config is None:
+ print("[ERROR] Follower arm 'main' configuration not found.")
+ return
+ motors_bus = FeetechMotorsBus(motor_config)
+ motors_bus.connect()
+
+ # Calibrate the follower arm.
+ calibrate_follower_arm(motors_bus, robot_config.calibration_dir)
+
+ # Create the LeKiwi robot instance.
+ robot = LeKiwi(motors_bus)
+
+ # Define the expected arm motor IDs.
+ arm_motor_ids = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
+
+ # Disable torque for each arm motor.
+ for motor in arm_motor_ids:
+ motors_bus.write("Torque_Enable", TorqueMode.DISABLED.value, motor)
+
+ # Set up ZeroMQ sockets.
+ context, cmd_socket, video_socket = setup_zmq_sockets(robot_config)
+
+ # Start the camera capture thread.
+ latest_images_dict = {}
+ images_lock = threading.Lock()
+ stop_event = threading.Event()
+ cam_thread = threading.Thread(
+ target=run_camera_capture, args=(cameras, images_lock, latest_images_dict, stop_event), daemon=True
+ )
+ cam_thread.start()
+
+ last_cmd_time = time.time()
+ print("LeKiwi robot server started. Waiting for commands...")
+
+ try:
+ while True:
+ loop_start_time = time.time()
+
+ # Process incoming commands (non-blocking).
+ while True:
+ try:
+ msg = cmd_socket.recv_string(zmq.NOBLOCK)
+ except zmq.Again:
+ break
+ try:
+ data = json.loads(msg)
+ # Process arm position commands.
+ if "arm_positions" in data:
+ arm_positions = data["arm_positions"]
+ if not isinstance(arm_positions, list):
+ print(f"[ERROR] Invalid arm_positions: {arm_positions}")
+ elif len(arm_positions) < len(arm_motor_ids):
+ print(
+ f"[WARNING] Received {len(arm_positions)} arm positions, expected {len(arm_motor_ids)}"
+ )
+ else:
+ for motor, pos in zip(arm_motor_ids, arm_positions, strict=False):
+ motors_bus.write("Goal_Position", pos, motor)
+ # Process wheel (base) commands.
+ if "raw_velocity" in data:
+ raw_command = data["raw_velocity"]
+ # Expect keys: "left_wheel", "back_wheel", "right_wheel".
+ command_speeds = [
+ int(raw_command.get("left_wheel", 0)),
+ int(raw_command.get("back_wheel", 0)),
+ int(raw_command.get("right_wheel", 0)),
+ ]
+ robot.set_velocity(command_speeds)
+ last_cmd_time = time.time()
+ except Exception as e:
+ print(f"[ERROR] Parsing message failed: {e}")
+
+ # Watchdog: stop the robot if no command is received for over 0.5 seconds.
+ now = time.time()
+ if now - last_cmd_time > 0.5:
+ robot.stop()
+ last_cmd_time = now
+
+ # Read current wheel speeds from the robot.
+ current_velocity = robot.read_velocity()
+
+ # Read the follower arm state from the motors bus.
+ follower_arm_state = []
+ for motor in arm_motor_ids:
+ try:
+ pos = motors_bus.read("Present_Position", motor)
+ # Convert the position to a float (or use as is if already numeric).
+ follower_arm_state.append(float(pos) if not isinstance(pos, (int, float)) else pos)
+ except Exception as e:
+ print(f"[ERROR] Reading motor {motor} failed: {e}")
+
+ # Get the latest camera images.
+ with images_lock:
+ images_dict_copy = dict(latest_images_dict)
+
+ # Build the observation dictionary.
+ observation = {
+ "images": images_dict_copy,
+ "present_speed": current_velocity,
+ "follower_arm_state": follower_arm_state,
+ }
+ # Send the observation over the video socket.
+ video_socket.send_string(json.dumps(observation))
+
+ # Ensure a short sleep to avoid overloading the CPU.
+ elapsed = time.time() - loop_start_time
+ time.sleep(
+ max(0.033 - elapsed, 0)
+ ) # If robot jitters increase the sleep and monitor cpu load with `top` in cmd
+ except KeyboardInterrupt:
+ print("Shutting down LeKiwi server.")
+ finally:
+ stop_event.set()
+ cam_thread.join()
+ robot.stop()
+ motors_bus.disconnect()
+ cmd_socket.close()
+ video_socket.close()
+ context.term()
diff --git a/lerobot/common/robot_devices/robots/manipulator.py b/lerobot/common/robot_devices/robots/manipulator.py
index e7f7cbb1..9173abc6 100644
--- a/lerobot/common/robot_devices/robots/manipulator.py
+++ b/lerobot/common/robot_devices/robots/manipulator.py
@@ -1,3 +1,17 @@
+# 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.
+
"""Contains logic to instantiate a robot, read information from its motors and cameras,
and send orders to its motors.
"""
@@ -44,7 +58,7 @@ class ManipulatorRobot:
# TODO(rcadene): Implement force feedback
"""This class allows to control any manipulator robot of various number of motors.
- Non exaustive list of robots:
+ Non exhaustive list of robots:
- [Koch v1.0](https://github.com/AlexanderKoch-Koch/low_cost_robot), with and without the wrist-to-elbow expansion, developed
by Alexander Koch from [Tau Robotics](https://tau-robotics.com)
- [Koch v1.1](https://github.com/jess-moss/koch-v1-1) developed by Jess Moss
@@ -55,7 +69,7 @@ class ManipulatorRobot:
robot = ManipulatorRobot(KochRobotConfig())
```
- Example of overwritting motors during instantiation:
+ Example of overwriting motors during instantiation:
```python
# Defines how to communicate with the motors of the leader and follower arms
leader_arms = {
@@ -90,7 +104,7 @@ class ManipulatorRobot:
robot = ManipulatorRobot(robot_config)
```
- Example of overwritting cameras during instantiation:
+ Example of overwriting cameras during instantiation:
```python
# Defines how to communicate with 2 cameras connected to the computer.
# Here, the webcam of the laptop and the phone (connected in USB to the laptop)
@@ -229,7 +243,7 @@ class ManipulatorRobot:
if self.robot_type in ["koch", "koch_bimanual", "aloha"]:
from lerobot.common.robot_devices.motors.dynamixel import TorqueMode
- elif self.robot_type in ["so100", "moss"]:
+ elif self.robot_type in ["so100", "moss", "lekiwi"]:
from lerobot.common.robot_devices.motors.feetech import TorqueMode
# We assume that at connection time, arms are in a rest position, and torque can
@@ -246,7 +260,7 @@ class ManipulatorRobot:
self.set_koch_robot_preset()
elif self.robot_type == "aloha":
self.set_aloha_robot_preset()
- elif self.robot_type in ["so100", "moss"]:
+ elif self.robot_type in ["so100", "moss", "lekiwi"]:
self.set_so100_robot_preset()
# Enable torque on all motors of the follower arms
@@ -299,7 +313,7 @@ class ManipulatorRobot:
calibration = run_arm_calibration(arm, self.robot_type, name, arm_type)
- elif self.robot_type in ["so100", "moss"]:
+ elif self.robot_type in ["so100", "moss", "lekiwi"]:
from lerobot.common.robot_devices.robots.feetech_calibration import (
run_arm_manual_calibration,
)
@@ -348,7 +362,7 @@ class ManipulatorRobot:
set_operating_mode_(self.follower_arms[name])
# Set better PID values to close the gap between recorded states and actions
- # TODO(rcadene): Implement an automatic procedure to set optimial PID values for each motor
+ # TODO(rcadene): Implement an automatic procedure to set optimal PID values for each motor
self.follower_arms[name].write("Position_P_Gain", 1500, "elbow_flex")
self.follower_arms[name].write("Position_I_Gain", 0, "elbow_flex")
self.follower_arms[name].write("Position_D_Gain", 600, "elbow_flex")
@@ -460,7 +474,7 @@ class ManipulatorRobot:
# Used when record_data=True
follower_goal_pos[name] = goal_pos
- goal_pos = goal_pos.numpy().astype(np.int32)
+ goal_pos = goal_pos.numpy().astype(np.float32)
self.follower_arms[name].write("Goal_Position", goal_pos)
self.logs[f"write_follower_{name}_goal_pos_dt_s"] = time.perf_counter() - before_fwrite_t
@@ -500,7 +514,7 @@ class ManipulatorRobot:
self.logs[f"read_camera_{name}_dt_s"] = self.cameras[name].logs["delta_timestamp_s"]
self.logs[f"async_read_camera_{name}_dt_s"] = time.perf_counter() - before_camread_t
- # Populate output dictionnaries
+ # Populate output dictionaries
obs_dict, action_dict = {}, {}
obs_dict["observation.state"] = state
action_dict["action"] = action
@@ -540,7 +554,7 @@ class ManipulatorRobot:
self.logs[f"read_camera_{name}_dt_s"] = self.cameras[name].logs["delta_timestamp_s"]
self.logs[f"async_read_camera_{name}_dt_s"] = time.perf_counter() - before_camread_t
- # Populate output dictionnaries and format to pytorch
+ # Populate output dictionaries and format to pytorch
obs_dict = {}
obs_dict["observation.state"] = state
for name in self.cameras:
@@ -582,7 +596,7 @@ class ManipulatorRobot:
action_sent.append(goal_pos)
# Send goal position to each follower
- goal_pos = goal_pos.numpy().astype(np.int32)
+ goal_pos = goal_pos.numpy().astype(np.float32)
self.follower_arms[name].write("Goal_Position", goal_pos)
return torch.cat(action_sent)
diff --git a/lerobot/common/robot_devices/robots/mobile_manipulator.py b/lerobot/common/robot_devices/robots/mobile_manipulator.py
new file mode 100644
index 00000000..385e218b
--- /dev/null
+++ b/lerobot/common/robot_devices/robots/mobile_manipulator.py
@@ -0,0 +1,703 @@
+# 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 base64
+import json
+import os
+import sys
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+import zmq
+
+from lerobot.common.robot_devices.cameras.utils import make_cameras_from_configs
+from lerobot.common.robot_devices.motors.feetech import TorqueMode
+from lerobot.common.robot_devices.motors.utils import MotorsBus, make_motors_buses_from_configs
+from lerobot.common.robot_devices.robots.configs import LeKiwiRobotConfig
+from lerobot.common.robot_devices.robots.feetech_calibration import run_arm_manual_calibration
+from lerobot.common.robot_devices.robots.utils import get_arm_id
+from lerobot.common.robot_devices.utils import RobotDeviceNotConnectedError
+
+PYNPUT_AVAILABLE = True
+try:
+ # Only import if there's a valid X server or if we're not on a Pi
+ if ("DISPLAY" not in os.environ) and ("linux" in sys.platform):
+ print("No DISPLAY set. Skipping pynput import.")
+ raise ImportError("pynput blocked intentionally due to no display.")
+
+ from pynput import keyboard
+except ImportError:
+ keyboard = None
+ PYNPUT_AVAILABLE = False
+except Exception as e:
+ keyboard = None
+ PYNPUT_AVAILABLE = False
+ print(f"Could not import pynput: {e}")
+
+
+class MobileManipulator:
+ """
+ MobileManipulator is a class for connecting to and controlling a remote mobile manipulator robot.
+ The robot includes a three omniwheel mobile base and a remote follower arm.
+ The leader arm is connected locally (on the laptop) and its joint positions are recorded and then
+ forwarded to the remote follower arm (after applying a safety clamp).
+ In parallel, keyboard teleoperation is used to generate raw velocity commands for the wheels.
+ """
+
+ def __init__(self, config: LeKiwiRobotConfig):
+ """
+ Expected keys in config:
+ - ip, port, video_port for the remote connection.
+ - calibration_dir, leader_arms, follower_arms, max_relative_target, etc.
+ """
+ self.robot_type = config.type
+ self.config = config
+ self.remote_ip = config.ip
+ self.remote_port = config.port
+ self.remote_port_video = config.video_port
+ self.calibration_dir = Path(self.config.calibration_dir)
+ self.logs = {}
+
+ self.teleop_keys = self.config.teleop_keys
+
+ # For teleoperation, the leader arm (local) is used to record the desired arm pose.
+ self.leader_arms = make_motors_buses_from_configs(self.config.leader_arms)
+
+ self.follower_arms = make_motors_buses_from_configs(self.config.follower_arms)
+
+ self.cameras = make_cameras_from_configs(self.config.cameras)
+
+ self.is_connected = False
+
+ self.last_frames = {}
+ self.last_present_speed = {}
+ self.last_remote_arm_state = torch.zeros(6, dtype=torch.float32)
+
+ # Define three speed levels and a current index
+ self.speed_levels = [
+ {"xy": 0.1, "theta": 30}, # slow
+ {"xy": 0.2, "theta": 60}, # medium
+ {"xy": 0.3, "theta": 90}, # fast
+ ]
+ self.speed_index = 0 # Start at slow
+
+ # ZeroMQ context and sockets.
+ self.context = None
+ self.cmd_socket = None
+ self.video_socket = None
+
+ # Keyboard state for base teleoperation.
+ self.running = True
+ self.pressed_keys = {
+ "forward": False,
+ "backward": False,
+ "left": False,
+ "right": False,
+ "rotate_left": False,
+ "rotate_right": False,
+ }
+
+ if PYNPUT_AVAILABLE:
+ print("pynput is available - enabling local keyboard listener.")
+ self.listener = keyboard.Listener(
+ on_press=self.on_press,
+ on_release=self.on_release,
+ )
+ self.listener.start()
+ else:
+ print("pynput not available - skipping local keyboard listener.")
+ self.listener = None
+
+ def get_motor_names(self, arms: dict[str, MotorsBus]) -> list:
+ return [f"{arm}_{motor}" for arm, bus in arms.items() for motor in bus.motors]
+
+ @property
+ def camera_features(self) -> dict:
+ cam_ft = {}
+ for cam_key, cam in self.cameras.items():
+ key = f"observation.images.{cam_key}"
+ cam_ft[key] = {
+ "shape": (cam.height, cam.width, cam.channels),
+ "names": ["height", "width", "channels"],
+ "info": None,
+ }
+ return cam_ft
+
+ @property
+ def motor_features(self) -> dict:
+ follower_arm_names = [
+ "shoulder_pan",
+ "shoulder_lift",
+ "elbow_flex",
+ "wrist_flex",
+ "wrist_roll",
+ "gripper",
+ ]
+ observations = ["x_mm", "y_mm", "theta"]
+ combined_names = follower_arm_names + observations
+ return {
+ "action": {
+ "dtype": "float32",
+ "shape": (len(combined_names),),
+ "names": combined_names,
+ },
+ "observation.state": {
+ "dtype": "float32",
+ "shape": (len(combined_names),),
+ "names": combined_names,
+ },
+ }
+
+ @property
+ def features(self):
+ return {**self.motor_features, **self.camera_features}
+
+ @property
+ def has_camera(self):
+ return len(self.cameras) > 0
+
+ @property
+ def num_cameras(self):
+ return len(self.cameras)
+
+ @property
+ def available_arms(self):
+ available = []
+ for name in self.leader_arms:
+ available.append(get_arm_id(name, "leader"))
+ for name in self.follower_arms:
+ available.append(get_arm_id(name, "follower"))
+ return available
+
+ def on_press(self, key):
+ try:
+ # Movement
+ if key.char == self.teleop_keys["forward"]:
+ self.pressed_keys["forward"] = True
+ elif key.char == self.teleop_keys["backward"]:
+ self.pressed_keys["backward"] = True
+ elif key.char == self.teleop_keys["left"]:
+ self.pressed_keys["left"] = True
+ elif key.char == self.teleop_keys["right"]:
+ self.pressed_keys["right"] = True
+ elif key.char == self.teleop_keys["rotate_left"]:
+ self.pressed_keys["rotate_left"] = True
+ elif key.char == self.teleop_keys["rotate_right"]:
+ self.pressed_keys["rotate_right"] = True
+
+ # Quit teleoperation
+ elif key.char == self.teleop_keys["quit"]:
+ self.running = False
+ return False
+
+ # Speed control
+ elif key.char == self.teleop_keys["speed_up"]:
+ self.speed_index = min(self.speed_index + 1, 2)
+ print(f"Speed index increased to {self.speed_index}")
+ elif key.char == self.teleop_keys["speed_down"]:
+ self.speed_index = max(self.speed_index - 1, 0)
+ print(f"Speed index decreased to {self.speed_index}")
+
+ except AttributeError:
+ # e.g., if key is special like Key.esc
+ if key == keyboard.Key.esc:
+ self.running = False
+ return False
+
+ def on_release(self, key):
+ try:
+ if hasattr(key, "char"):
+ if key.char == self.teleop_keys["forward"]:
+ self.pressed_keys["forward"] = False
+ elif key.char == self.teleop_keys["backward"]:
+ self.pressed_keys["backward"] = False
+ elif key.char == self.teleop_keys["left"]:
+ self.pressed_keys["left"] = False
+ elif key.char == self.teleop_keys["right"]:
+ self.pressed_keys["right"] = False
+ elif key.char == self.teleop_keys["rotate_left"]:
+ self.pressed_keys["rotate_left"] = False
+ elif key.char == self.teleop_keys["rotate_right"]:
+ self.pressed_keys["rotate_right"] = False
+ except AttributeError:
+ pass
+
+ def connect(self):
+ if not self.leader_arms:
+ raise ValueError("MobileManipulator has no leader arm to connect.")
+ for name in self.leader_arms:
+ print(f"Connecting {name} leader arm.")
+ self.calibrate_leader()
+
+ # Set up ZeroMQ sockets to communicate with the remote mobile robot.
+ self.context = zmq.Context()
+ self.cmd_socket = self.context.socket(zmq.PUSH)
+ connection_string = f"tcp://{self.remote_ip}:{self.remote_port}"
+ self.cmd_socket.connect(connection_string)
+ self.cmd_socket.setsockopt(zmq.CONFLATE, 1)
+ self.video_socket = self.context.socket(zmq.PULL)
+ video_connection = f"tcp://{self.remote_ip}:{self.remote_port_video}"
+ self.video_socket.connect(video_connection)
+ self.video_socket.setsockopt(zmq.CONFLATE, 1)
+ print(
+ f"[INFO] Connected to remote robot at {connection_string} and video stream at {video_connection}."
+ )
+ self.is_connected = True
+
+ def load_or_run_calibration_(self, name, arm, arm_type):
+ arm_id = get_arm_id(name, arm_type)
+ arm_calib_path = self.calibration_dir / f"{arm_id}.json"
+
+ if arm_calib_path.exists():
+ with open(arm_calib_path) as f:
+ calibration = json.load(f)
+ else:
+ print(f"Missing calibration file '{arm_calib_path}'")
+ calibration = run_arm_manual_calibration(arm, self.robot_type, name, arm_type)
+ print(f"Calibration is done! Saving calibration file '{arm_calib_path}'")
+ arm_calib_path.parent.mkdir(parents=True, exist_ok=True)
+ with open(arm_calib_path, "w") as f:
+ json.dump(calibration, f)
+
+ return calibration
+
+ def calibrate_leader(self):
+ for name, arm in self.leader_arms.items():
+ # Connect the bus
+ arm.connect()
+
+ # Disable torque on all motors
+ for motor_id in arm.motors:
+ arm.write("Torque_Enable", TorqueMode.DISABLED.value, motor_id)
+
+ # Now run calibration
+ calibration = self.load_or_run_calibration_(name, arm, "leader")
+ arm.set_calibration(calibration)
+
+ def calibrate_follower(self):
+ for name, bus in self.follower_arms.items():
+ bus.connect()
+
+ # Disable torque on all motors
+ for motor_id in bus.motors:
+ bus.write("Torque_Enable", 0, motor_id)
+
+ # Then filter out wheels
+ arm_only_dict = {k: v for k, v in bus.motors.items() if not k.startswith("wheel_")}
+ if not arm_only_dict:
+ continue
+
+ original_motors = bus.motors
+ bus.motors = arm_only_dict
+
+ calibration = self.load_or_run_calibration_(name, bus, "follower")
+ bus.set_calibration(calibration)
+
+ bus.motors = original_motors
+
+ def _get_data(self):
+ """
+ Polls the video socket for up to 15 ms. If data arrives, decode only
+ the *latest* message, returning frames, speed, and arm state. If
+ nothing arrives for any field, use the last known values.
+ """
+ frames = {}
+ present_speed = {}
+ remote_arm_state_tensor = torch.zeros(6, dtype=torch.float32)
+
+ # Poll up to 15 ms
+ poller = zmq.Poller()
+ poller.register(self.video_socket, zmq.POLLIN)
+ socks = dict(poller.poll(15))
+ if self.video_socket not in socks or socks[self.video_socket] != zmq.POLLIN:
+ # No new data arrived → reuse ALL old data
+ return (self.last_frames, self.last_present_speed, self.last_remote_arm_state)
+
+ # Drain all messages, keep only the last
+ last_msg = None
+ while True:
+ try:
+ obs_string = self.video_socket.recv_string(zmq.NOBLOCK)
+ last_msg = obs_string
+ except zmq.Again:
+ break
+
+ if not last_msg:
+ # No new message → also reuse old
+ return (self.last_frames, self.last_present_speed, self.last_remote_arm_state)
+
+ # Decode only the final message
+ try:
+ observation = json.loads(last_msg)
+
+ images_dict = observation.get("images", {})
+ new_speed = observation.get("present_speed", {})
+ new_arm_state = observation.get("follower_arm_state", None)
+
+ # Convert images
+ for cam_name, image_b64 in images_dict.items():
+ if image_b64:
+ jpg_data = base64.b64decode(image_b64)
+ np_arr = np.frombuffer(jpg_data, dtype=np.uint8)
+ frame_candidate = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+ if frame_candidate is not None:
+ frames[cam_name] = frame_candidate
+
+ # If remote_arm_state is None and frames is None there is no message then use the previous message
+ if new_arm_state is not None and frames is not None:
+ self.last_frames = frames
+
+ remote_arm_state_tensor = torch.tensor(new_arm_state, dtype=torch.float32)
+ self.last_remote_arm_state = remote_arm_state_tensor
+
+ present_speed = new_speed
+ self.last_present_speed = new_speed
+ else:
+ frames = self.last_frames
+
+ remote_arm_state_tensor = self.last_remote_arm_state
+
+ present_speed = self.last_present_speed
+
+ except Exception as e:
+ print(f"[DEBUG] Error decoding video message: {e}")
+ # If decode fails, fall back to old data
+ return (self.last_frames, self.last_present_speed, self.last_remote_arm_state)
+
+ return frames, present_speed, remote_arm_state_tensor
+
+ def _process_present_speed(self, present_speed: dict) -> torch.Tensor:
+ state_tensor = torch.zeros(3, dtype=torch.int32)
+ if present_speed:
+ decoded = {key: MobileManipulator.raw_to_degps(value) for key, value in present_speed.items()}
+ if "1" in decoded:
+ state_tensor[0] = decoded["1"]
+ if "2" in decoded:
+ state_tensor[1] = decoded["2"]
+ if "3" in decoded:
+ state_tensor[2] = decoded["3"]
+ return state_tensor
+
+ def teleop_step(
+ self, record_data: bool = False
+ ) -> None | tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
+ if not self.is_connected:
+ raise RobotDeviceNotConnectedError("MobileManipulator is not connected. Run `connect()` first.")
+
+ speed_setting = self.speed_levels[self.speed_index]
+ xy_speed = speed_setting["xy"] # e.g. 0.1, 0.25, or 0.4
+ theta_speed = speed_setting["theta"] # e.g. 30, 60, or 90
+
+ # Prepare to assign the position of the leader to the follower
+ arm_positions = []
+ for name in self.leader_arms:
+ pos = self.leader_arms[name].read("Present_Position")
+ pos_tensor = torch.from_numpy(pos).float()
+ arm_positions.extend(pos_tensor.tolist())
+
+ 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"]:
+ y_cmd += xy_speed
+ 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"]:
+ theta_cmd -= theta_speed
+
+ wheel_commands = self.body_to_wheel_raw(x_cmd, y_cmd, theta_cmd)
+
+ message = {"raw_velocity": wheel_commands, "arm_positions": arm_positions}
+ self.cmd_socket.send_string(json.dumps(message))
+
+ if not record_data:
+ return
+
+ obs_dict = self.capture_observation()
+
+ arm_state_tensor = torch.tensor(arm_positions, dtype=torch.float32)
+
+ wheel_velocity_tuple = self.wheel_raw_to_body(wheel_commands)
+ wheel_velocity_mm = (
+ wheel_velocity_tuple[0] * 1000.0,
+ wheel_velocity_tuple[1] * 1000.0,
+ wheel_velocity_tuple[2],
+ )
+ wheel_tensor = torch.tensor(wheel_velocity_mm, dtype=torch.float32)
+ action_tensor = torch.cat([arm_state_tensor, wheel_tensor])
+ action_dict = {"action": action_tensor}
+
+ return obs_dict, action_dict
+
+ def capture_observation(self) -> dict:
+ """
+ Capture observations from the remote robot: current follower arm positions,
+ present wheel speeds (converted to body-frame velocities: x, y, theta),
+ and a camera frame.
+ """
+ if not self.is_connected:
+ raise RobotDeviceNotConnectedError("Not connected. Run `connect()` first.")
+
+ frames, present_speed, remote_arm_state_tensor = self._get_data()
+
+ body_state = self.wheel_raw_to_body(present_speed)
+
+ body_state_mm = (body_state[0] * 1000.0, body_state[1] * 1000.0, body_state[2]) # Convert x,y to mm/s
+ wheel_state_tensor = torch.tensor(body_state_mm, dtype=torch.float32)
+ combined_state_tensor = torch.cat((remote_arm_state_tensor, wheel_state_tensor), dim=0)
+
+ obs_dict = {"observation.state": combined_state_tensor}
+
+ # Loop over each configured camera
+ for cam_name, cam in self.cameras.items():
+ frame = frames.get(cam_name, None)
+ if frame is None:
+ # Create a black image using the camera's configured width, height, and channels
+ frame = np.zeros((cam.height, cam.width, cam.channels), dtype=np.uint8)
+ obs_dict[f"observation.images.{cam_name}"] = torch.from_numpy(frame)
+
+ return obs_dict
+
+ def send_action(self, action: torch.Tensor) -> torch.Tensor:
+ if not self.is_connected:
+ raise RobotDeviceNotConnectedError("Not connected. Run `connect()` first.")
+
+ # Ensure the action tensor has at least 9 elements:
+ # - First 6: arm positions.
+ # - Last 3: base commands.
+ if action.numel() < 9:
+ # Pad with zeros if there are not enough elements.
+ padded = torch.zeros(9, dtype=action.dtype)
+ padded[: action.numel()] = action
+ action = padded
+
+ # Extract arm and base actions.
+ arm_actions = action[:6].flatten()
+ base_actions = action[6:].flatten()
+
+ x_cmd_mm = base_actions[0].item() # mm/s
+ y_cmd_mm = base_actions[1].item() # mm/s
+ theta_cmd = base_actions[2].item() # deg/s
+
+ # Convert mm/s to m/s for the kinematics calculations.
+ x_cmd = x_cmd_mm / 1000.0 # m/s
+ y_cmd = y_cmd_mm / 1000.0 # m/s
+
+ # Compute wheel commands from body commands.
+ wheel_commands = self.body_to_wheel_raw(x_cmd, y_cmd, theta_cmd)
+
+ arm_positions_list = arm_actions.tolist()
+
+ message = {"raw_velocity": wheel_commands, "arm_positions": arm_positions_list}
+ self.cmd_socket.send_string(json.dumps(message))
+
+ return action
+
+ def print_logs(self):
+ pass
+
+ def disconnect(self):
+ if not self.is_connected:
+ raise RobotDeviceNotConnectedError("Not connected.")
+ if self.cmd_socket:
+ stop_cmd = {
+ "raw_velocity": {"left_wheel": 0, "back_wheel": 0, "right_wheel": 0},
+ "arm_positions": {},
+ }
+ self.cmd_socket.send_string(json.dumps(stop_cmd))
+ self.cmd_socket.close()
+ if self.video_socket:
+ self.video_socket.close()
+ if self.context:
+ self.context.term()
+ if PYNPUT_AVAILABLE:
+ self.listener.stop()
+ self.is_connected = False
+ print("[INFO] Disconnected from remote robot.")
+
+ def __del__(self):
+ if getattr(self, "is_connected", False):
+ self.disconnect()
+ if PYNPUT_AVAILABLE:
+ self.listener.stop()
+
+ @staticmethod
+ def degps_to_raw(degps: float) -> int:
+ steps_per_deg = 4096.0 / 360.0
+ speed_in_steps = abs(degps) * steps_per_deg
+ speed_int = int(round(speed_in_steps))
+ if speed_int > 0x7FFF:
+ speed_int = 0x7FFF
+ if degps < 0:
+ return speed_int | 0x8000
+ else:
+ return speed_int & 0x7FFF
+
+ @staticmethod
+ def raw_to_degps(raw_speed: int) -> float:
+ steps_per_deg = 4096.0 / 360.0
+ magnitude = raw_speed & 0x7FFF
+ degps = magnitude / steps_per_deg
+ if raw_speed & 0x8000:
+ degps = -degps
+ return degps
+
+ def body_to_wheel_raw(
+ self,
+ x_cmd: float,
+ y_cmd: float,
+ theta_cmd: float,
+ wheel_radius: float = 0.05,
+ base_radius: float = 0.125,
+ max_raw: int = 3000,
+ ) -> dict:
+ """
+ Convert desired body-frame velocities into wheel raw commands.
+
+ Parameters:
+ x_cmd : Linear velocity in x (m/s).
+ y_cmd : Linear velocity in y (m/s).
+ theta_cmd : Rotational velocity (deg/s).
+ wheel_radius: Radius of each wheel (meters).
+ base_radius : Distance from the center of rotation to each wheel (meters).
+ max_raw : Maximum allowed raw command (ticks) per wheel.
+
+ Returns:
+ A dictionary with wheel raw commands:
+ {"left_wheel": value, "back_wheel": value, "right_wheel": value}.
+
+ Notes:
+ - Internally, the method converts theta_cmd to rad/s for the kinematics.
+ - The raw command is computed from the wheels angular speed in deg/s
+ using degps_to_raw(). If any command exceeds max_raw, all commands
+ are scaled down proportionally.
+ """
+ # Convert rotational velocity from deg/s to rad/s.
+ theta_rad = theta_cmd * (np.pi / 180.0)
+ # Create the body velocity vector [x, y, theta_rad].
+ velocity_vector = np.array([x_cmd, y_cmd, theta_rad])
+
+ # 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])
+
+ # Compute each wheel’s linear speed (m/s) and then its angular speed (rad/s).
+ wheel_linear_speeds = m.dot(velocity_vector)
+ wheel_angular_speeds = wheel_linear_speeds / wheel_radius
+
+ # Convert wheel angular speeds from rad/s to deg/s.
+ wheel_degps = wheel_angular_speeds * (180.0 / np.pi)
+
+ # Scaling
+ steps_per_deg = 4096.0 / 360.0
+ raw_floats = [abs(degps) * steps_per_deg for degps in wheel_degps]
+ max_raw_computed = max(raw_floats)
+ if max_raw_computed > max_raw:
+ scale = max_raw / max_raw_computed
+ wheel_degps = wheel_degps * scale
+
+ # Convert each wheel’s angular speed (deg/s) to a raw integer.
+ wheel_raw = [MobileManipulator.degps_to_raw(deg) for deg in wheel_degps]
+
+ return {"left_wheel": wheel_raw[0], "back_wheel": wheel_raw[1], "right_wheel": wheel_raw[2]}
+
+ def wheel_raw_to_body(
+ self, wheel_raw: dict, wheel_radius: float = 0.05, base_radius: float = 0.125
+ ) -> tuple:
+ """
+ Convert wheel raw command feedback back into body-frame velocities.
+
+ Parameters:
+ wheel_raw : Dictionary with raw wheel commands (keys: "left_wheel", "back_wheel", "right_wheel").
+ wheel_radius: Radius of each wheel (meters).
+ base_radius : Distance from the robot center to each wheel (meters).
+
+ Returns:
+ A tuple (x_cmd, y_cmd, theta_cmd) where:
+ x_cmd : Linear velocity in x (m/s).
+ y_cmd : Linear velocity in y (m/s).
+ theta_cmd : Rotational velocity in deg/s.
+ """
+ # Extract the raw values in order.
+ raw_list = [
+ int(wheel_raw.get("left_wheel", 0)),
+ int(wheel_raw.get("back_wheel", 0)),
+ int(wheel_raw.get("right_wheel", 0)),
+ ]
+
+ # Convert each raw command back to an angular speed in deg/s.
+ wheel_degps = np.array([MobileManipulator.raw_to_degps(r) for r in raw_list])
+ # Convert from deg/s to rad/s.
+ wheel_radps = wheel_degps * (np.pi / 180.0)
+ # Compute each wheel’s linear speed (m/s) from its angular speed.
+ wheel_linear_speeds = wheel_radps * wheel_radius
+
+ # 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.
+ m_inv = np.linalg.inv(m)
+ velocity_vector = m_inv.dot(wheel_linear_speeds)
+ x_cmd, y_cmd, theta_rad = velocity_vector
+ theta_cmd = theta_rad * (180.0 / np.pi)
+ return (x_cmd, y_cmd, theta_cmd)
+
+
+class LeKiwi:
+ def __init__(self, motor_bus):
+ """
+ Initializes the LeKiwi with Feetech motors bus.
+ """
+ self.motor_bus = motor_bus
+ self.motor_ids = ["left_wheel", "back_wheel", "right_wheel"]
+
+ # Initialize motors in velocity mode.
+ self.motor_bus.write("Lock", 0)
+ self.motor_bus.write("Mode", [1, 1, 1], self.motor_ids)
+ self.motor_bus.write("Lock", 1)
+ print("Motors set to velocity mode.")
+
+ def read_velocity(self):
+ """
+ Reads the raw speeds for all wheels. Returns a dictionary with motor names:
+ """
+ raw_speeds = self.motor_bus.read("Present_Speed", self.motor_ids)
+ return {
+ "left_wheel": int(raw_speeds[0]),
+ "back_wheel": int(raw_speeds[1]),
+ "right_wheel": int(raw_speeds[2]),
+ }
+
+ def set_velocity(self, command_speeds):
+ """
+ Sends raw velocity commands (16-bit encoded values) directly to the motor bus.
+ The order of speeds must correspond to self.motor_ids.
+ """
+ self.motor_bus.write("Goal_Speed", command_speeds, self.motor_ids)
+
+ def stop(self):
+ """Stops the robot by setting all motor speeds to zero."""
+ self.motor_bus.write("Goal_Speed", [0, 0, 0], self.motor_ids)
+ print("Motors stopped.")
diff --git a/lerobot/common/robot_devices/robots/stretch.py b/lerobot/common/robot_devices/robots/stretch.py
index b63bf941..9cfe6e49 100644
--- a/lerobot/common/robot_devices/robots/stretch.py
+++ b/lerobot/common/robot_devices/robots/stretch.py
@@ -108,7 +108,7 @@ class StretchRobot(StretchAPI):
self.logs[f"read_camera_{name}_dt_s"] = self.cameras[name].logs["delta_timestamp_s"]
self.logs[f"async_read_camera_{name}_dt_s"] = time.perf_counter() - before_camread_t
- # Populate output dictionnaries
+ # Populate output dictionaries
obs_dict, action_dict = {}, {}
obs_dict["observation.state"] = state
action_dict["action"] = action
@@ -153,7 +153,7 @@ class StretchRobot(StretchAPI):
self.logs[f"read_camera_{name}_dt_s"] = self.cameras[name].logs["delta_timestamp_s"]
self.logs[f"async_read_camera_{name}_dt_s"] = time.perf_counter() - before_camread_t
- # Populate output dictionnaries
+ # Populate output dictionaries
obs_dict = {}
obs_dict["observation.state"] = state
for name in self.cameras:
diff --git a/lerobot/common/robot_devices/robots/utils.py b/lerobot/common/robot_devices/robots/utils.py
index 86ff6473..dab514d5 100644
--- a/lerobot/common/robot_devices/robots/utils.py
+++ b/lerobot/common/robot_devices/robots/utils.py
@@ -1,9 +1,24 @@
+# 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.
+
from typing import Protocol
from lerobot.common.robot_devices.robots.configs import (
AlohaRobotConfig,
KochBimanualRobotConfig,
KochRobotConfig,
+ LeKiwiRobotConfig,
ManipulatorRobotConfig,
MossRobotConfig,
RobotConfig,
@@ -45,6 +60,8 @@ def make_robot_config(robot_type: str, **kwargs) -> RobotConfig:
return So100RobotConfig(**kwargs)
elif robot_type == "stretch":
return StretchRobotConfig(**kwargs)
+ elif robot_type == "lekiwi":
+ return LeKiwiRobotConfig(**kwargs)
else:
raise ValueError(f"Robot type '{robot_type}' is not available.")
@@ -54,6 +71,10 @@ def make_robot_from_config(config: RobotConfig):
from lerobot.common.robot_devices.robots.manipulator import ManipulatorRobot
return ManipulatorRobot(config)
+ elif isinstance(config, LeKiwiRobotConfig):
+ from lerobot.common.robot_devices.robots.mobile_manipulator import MobileManipulator
+
+ return MobileManipulator(config)
else:
from lerobot.common.robot_devices.robots.stretch import StretchRobot
diff --git a/lerobot/common/robot_devices/utils.py b/lerobot/common/robot_devices/utils.py
index 19bb637e..837c9d2e 100644
--- a/lerobot/common/robot_devices/utils.py
+++ b/lerobot/common/robot_devices/utils.py
@@ -1,3 +1,17 @@
+# 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 platform
import time
diff --git a/lerobot/common/utils/hub.py b/lerobot/common/utils/hub.py
index 63fcf918..df7435c0 100644
--- a/lerobot/common/utils/hub.py
+++ b/lerobot/common/utils/hub.py
@@ -1,3 +1,17 @@
+# 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.
+
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Any, Type, TypeVar
diff --git a/lerobot/common/utils/io_utils.py b/lerobot/common/utils/io_utils.py
index b85f17c7..da0be1c7 100644
--- a/lerobot/common/utils/io_utils.py
+++ b/lerobot/common/utils/io_utils.py
@@ -13,10 +13,16 @@
# 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 json
import warnings
+from pathlib import Path
+from typing import TypeVar
import imageio
+JsonLike = str | int | float | bool | None | list["JsonLike"] | dict[str, "JsonLike"] | tuple["JsonLike", ...]
+T = TypeVar("T", bound=JsonLike)
+
def write_video(video_path, stacked_frames, fps):
# Filter out DeprecationWarnings raised from pkg_resources
@@ -25,3 +31,81 @@ def write_video(video_path, stacked_frames, fps):
"ignore", "pkg_resources is deprecated as an API", category=DeprecationWarning
)
imageio.mimsave(video_path, stacked_frames, fps=fps)
+
+
+def deserialize_json_into_object(fpath: Path, obj: T) -> T:
+ """
+ Loads the JSON data from `fpath` and recursively fills `obj` with the
+ corresponding values (strictly matching structure and types).
+ Tuples in `obj` are expected to be lists in the JSON data, which will be
+ converted back into tuples.
+ """
+ with open(fpath, encoding="utf-8") as f:
+ data = json.load(f)
+
+ def _deserialize(target, source):
+ """
+ Recursively overwrite the structure in `target` with data from `source`,
+ performing strict checks on structure and type.
+ Returns the updated version of `target` (especially important for tuples).
+ """
+
+ # If the target is a dictionary, source must be a dictionary as well.
+ if isinstance(target, dict):
+ if not isinstance(source, dict):
+ raise TypeError(f"Type mismatch: expected dict, got {type(source)}")
+
+ # Check that they have exactly the same set of keys.
+ if target.keys() != source.keys():
+ raise ValueError(
+ f"Dictionary keys do not match.\nExpected: {target.keys()}, got: {source.keys()}"
+ )
+
+ # Recursively update each key.
+ for k in target:
+ target[k] = _deserialize(target[k], source[k])
+
+ return target
+
+ # If the target is a list, source must be a list as well.
+ elif isinstance(target, list):
+ if not isinstance(source, list):
+ raise TypeError(f"Type mismatch: expected list, got {type(source)}")
+
+ # Check length
+ if len(target) != len(source):
+ raise ValueError(f"List length mismatch: expected {len(target)}, got {len(source)}")
+
+ # Recursively update each element.
+ for i in range(len(target)):
+ target[i] = _deserialize(target[i], source[i])
+
+ return target
+
+ # If the target is a tuple, the source must be a list in JSON,
+ # which we'll convert back to a tuple.
+ elif isinstance(target, tuple):
+ if not isinstance(source, list):
+ raise TypeError(f"Type mismatch: expected list (for tuple), got {type(source)}")
+
+ if len(target) != len(source):
+ raise ValueError(f"Tuple length mismatch: expected {len(target)}, got {len(source)}")
+
+ # Convert each element, forming a new tuple.
+ converted_items = []
+ for t_item, s_item in zip(target, source, strict=False):
+ converted_items.append(_deserialize(t_item, s_item))
+
+ # Return a brand new tuple (tuples are immutable in Python).
+ return tuple(converted_items)
+
+ # Otherwise, we're dealing with a "primitive" (int, float, str, bool, None).
+ else:
+ # Check the exact type. If these must match 1:1, do:
+ if type(target) is not type(source):
+ raise TypeError(f"Type mismatch: expected {type(target)}, got {type(source)}")
+ return source
+
+ # Perform the in-place/recursive deserialization
+ updated_obj = _deserialize(obj, data)
+ return updated_obj
diff --git a/lerobot/common/utils/logging_utils.py b/lerobot/common/utils/logging_utils.py
new file mode 100644
index 00000000..b99c348f
--- /dev/null
+++ b/lerobot/common/utils/logging_utils.py
@@ -0,0 +1,163 @@
+#!/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.
+from typing import Any
+
+from lerobot.common.utils.utils import format_big_number
+
+
+class AverageMeter:
+ """
+ Computes and stores the average and current value
+ Adapted from https://github.com/pytorch/examples/blob/main/imagenet/main.py
+ """
+
+ def __init__(self, name: str, fmt: str = ":f"):
+ self.name = name
+ self.fmt = fmt
+ self.reset()
+
+ def reset(self) -> None:
+ self.val = 0.0
+ self.avg = 0.0
+ self.sum = 0.0
+ self.count = 0.0
+
+ def update(self, val: float, n: int = 1) -> None:
+ self.val = val
+ self.sum += val * n
+ self.count += n
+ self.avg = self.sum / self.count
+
+ def __str__(self):
+ fmtstr = "{name}:{avg" + self.fmt + "}"
+ return fmtstr.format(**self.__dict__)
+
+
+class MetricsTracker:
+ """
+ A helper class to track and log metrics over time.
+
+ Usage pattern:
+
+ ```python
+ # initialize, potentially with non-zero initial step (e.g. if resuming run)
+ metrics = {"loss": AverageMeter("loss", ":.3f")}
+ train_metrics = MetricsTracker(cfg, dataset, metrics, initial_step=step)
+
+ # update metrics derived from step (samples, episodes, epochs) at each training step
+ train_metrics.step()
+
+ # update various metrics
+ loss = policy.forward(batch)
+ train_metrics.loss = loss
+
+ # display current metrics
+ logging.info(train_metrics)
+
+ # export for wandb
+ wandb.log(train_metrics.to_dict())
+
+ # reset averages after logging
+ train_metrics.reset_averages()
+ ```
+ """
+
+ __keys__ = [
+ "_batch_size",
+ "_num_frames",
+ "_avg_samples_per_ep",
+ "metrics",
+ "steps",
+ "samples",
+ "episodes",
+ "epochs",
+ ]
+
+ def __init__(
+ self,
+ batch_size: int,
+ num_frames: int,
+ num_episodes: int,
+ metrics: dict[str, AverageMeter],
+ initial_step: int = 0,
+ ):
+ self.__dict__.update({k: None for k in self.__keys__})
+ self._batch_size = batch_size
+ self._num_frames = num_frames
+ self._avg_samples_per_ep = num_frames / num_episodes
+ self.metrics = metrics
+
+ self.steps = initial_step
+ # A sample is an (observation,action) pair, where observation and action
+ # can be on multiple timestamps. In a batch, we have `batch_size` number of samples.
+ self.samples = self.steps * self._batch_size
+ self.episodes = self.samples / self._avg_samples_per_ep
+ self.epochs = self.samples / self._num_frames
+
+ def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any:
+ if name in self.__dict__:
+ return self.__dict__[name]
+ elif name in self.metrics:
+ return self.metrics[name]
+ else:
+ raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+
+ def __setattr__(self, name: str, value: Any) -> None:
+ if name in self.__dict__:
+ super().__setattr__(name, value)
+ elif name in self.metrics:
+ self.metrics[name].update(value)
+ else:
+ raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+
+ def step(self) -> None:
+ """
+ Updates metrics that depend on 'step' for one step.
+ """
+ self.steps += 1
+ self.samples += self._batch_size
+ self.episodes = self.samples / self._avg_samples_per_ep
+ self.epochs = self.samples / self._num_frames
+
+ def __str__(self) -> str:
+ display_list = [
+ f"step:{format_big_number(self.steps)}",
+ # number of samples seen during training
+ f"smpl:{format_big_number(self.samples)}",
+ # number of episodes seen during training
+ f"ep:{format_big_number(self.episodes)}",
+ # number of time all unique samples are seen
+ f"epch:{self.epochs:.2f}",
+ *[str(m) for m in self.metrics.values()],
+ ]
+ return " ".join(display_list)
+
+ def to_dict(self, use_avg: bool = True) -> dict[str, int | float]:
+ """
+ Returns the current metric values (or averages if `use_avg=True`) as a dict.
+ """
+ return {
+ "steps": self.steps,
+ "samples": self.samples,
+ "episodes": self.episodes,
+ "epochs": self.epochs,
+ **{k: m.avg if use_avg else m.val for k, m in self.metrics.items()},
+ }
+
+ def reset_averages(self) -> None:
+ """Resets average meters."""
+ for m in self.metrics.values():
+ m.reset()
diff --git a/lerobot/common/utils/random_utils.py b/lerobot/common/utils/random_utils.py
new file mode 100644
index 00000000..3d9bf4dd
--- /dev/null
+++ b/lerobot/common/utils/random_utils.py
@@ -0,0 +1,191 @@
+#!/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 random
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Generator
+
+import numpy as np
+import torch
+from safetensors.torch import load_file, save_file
+
+from lerobot.common.constants import RNG_STATE
+from lerobot.common.datasets.utils import flatten_dict, unflatten_dict
+
+
+def serialize_python_rng_state() -> dict[str, torch.Tensor]:
+ """
+ Returns the rng state for `random` in the form of a flat dict[str, torch.Tensor] to be saved using
+ `safetensors.save_file()` or `torch.save()`.
+ """
+ py_state = random.getstate()
+ return {
+ "py_rng_version": torch.tensor([py_state[0]], dtype=torch.int64),
+ "py_rng_state": torch.tensor(py_state[1], dtype=torch.int64),
+ }
+
+
+def deserialize_python_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None:
+ """
+ Restores the rng state for `random` from a dictionary produced by `serialize_python_rng_state()`.
+ """
+ py_state = (rng_state_dict["py_rng_version"].item(), tuple(rng_state_dict["py_rng_state"].tolist()), None)
+ random.setstate(py_state)
+
+
+def serialize_numpy_rng_state() -> dict[str, torch.Tensor]:
+ """
+ Returns the rng state for `numpy` in the form of a flat dict[str, torch.Tensor] to be saved using
+ `safetensors.save_file()` or `torch.save()`.
+ """
+ np_state = np.random.get_state()
+ # Ensure no breaking changes from numpy
+ assert np_state[0] == "MT19937"
+ return {
+ "np_rng_state_values": torch.tensor(np_state[1], dtype=torch.int64),
+ "np_rng_state_index": torch.tensor([np_state[2]], dtype=torch.int64),
+ "np_rng_has_gauss": torch.tensor([np_state[3]], dtype=torch.int64),
+ "np_rng_cached_gaussian": torch.tensor([np_state[4]], dtype=torch.float32),
+ }
+
+
+def deserialize_numpy_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None:
+ """
+ Restores the rng state for `numpy` from a dictionary produced by `serialize_numpy_rng_state()`.
+ """
+ np_state = (
+ "MT19937",
+ rng_state_dict["np_rng_state_values"].numpy(),
+ rng_state_dict["np_rng_state_index"].item(),
+ rng_state_dict["np_rng_has_gauss"].item(),
+ rng_state_dict["np_rng_cached_gaussian"].item(),
+ )
+ np.random.set_state(np_state)
+
+
+def serialize_torch_rng_state() -> dict[str, torch.Tensor]:
+ """
+ Returns the rng state for `torch` in the form of a flat dict[str, torch.Tensor] to be saved using
+ `safetensors.save_file()` or `torch.save()`.
+ """
+ torch_rng_state_dict = {"torch_rng_state": torch.get_rng_state()}
+ if torch.cuda.is_available():
+ torch_rng_state_dict["torch_cuda_rng_state"] = torch.cuda.get_rng_state()
+ return torch_rng_state_dict
+
+
+def deserialize_torch_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None:
+ """
+ Restores the rng state for `torch` from a dictionary produced by `serialize_torch_rng_state()`.
+ """
+ torch.set_rng_state(rng_state_dict["torch_rng_state"])
+ if torch.cuda.is_available() and "torch_cuda_rng_state" in rng_state_dict:
+ torch.cuda.set_rng_state(rng_state_dict["torch_cuda_rng_state"])
+
+
+def serialize_rng_state() -> dict[str, torch.Tensor]:
+ """
+ Returns the rng state for `random`, `numpy`, and `torch`, in the form of a flat
+ dict[str, torch.Tensor] to be saved using `safetensors.save_file()` `torch.save()`.
+ """
+ py_rng_state_dict = serialize_python_rng_state()
+ np_rng_state_dict = serialize_numpy_rng_state()
+ torch_rng_state_dict = serialize_torch_rng_state()
+
+ return {
+ **py_rng_state_dict,
+ **np_rng_state_dict,
+ **torch_rng_state_dict,
+ }
+
+
+def deserialize_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None:
+ """
+ Restores the rng state for `random`, `numpy`, and `torch` from a dictionary produced by
+ `serialize_rng_state()`.
+ """
+ py_rng_state_dict = {k: v for k, v in rng_state_dict.items() if k.startswith("py")}
+ np_rng_state_dict = {k: v for k, v in rng_state_dict.items() if k.startswith("np")}
+ torch_rng_state_dict = {k: v for k, v in rng_state_dict.items() if k.startswith("torch")}
+
+ deserialize_python_rng_state(py_rng_state_dict)
+ deserialize_numpy_rng_state(np_rng_state_dict)
+ deserialize_torch_rng_state(torch_rng_state_dict)
+
+
+def save_rng_state(save_dir: Path) -> None:
+ rng_state_dict = serialize_rng_state()
+ flat_rng_state_dict = flatten_dict(rng_state_dict)
+ save_file(flat_rng_state_dict, save_dir / RNG_STATE)
+
+
+def load_rng_state(save_dir: Path) -> None:
+ flat_rng_state_dict = load_file(save_dir / RNG_STATE)
+ rng_state_dict = unflatten_dict(flat_rng_state_dict)
+ deserialize_rng_state(rng_state_dict)
+
+
+def get_rng_state() -> dict[str, Any]:
+ """Get the random state for `random`, `numpy`, and `torch`."""
+ random_state_dict = {
+ "random_state": random.getstate(),
+ "numpy_random_state": np.random.get_state(),
+ "torch_random_state": torch.random.get_rng_state(),
+ }
+ if torch.cuda.is_available():
+ random_state_dict["torch_cuda_random_state"] = torch.cuda.random.get_rng_state()
+ return random_state_dict
+
+
+def set_rng_state(random_state_dict: dict[str, Any]):
+ """Set the random state for `random`, `numpy`, and `torch`.
+
+ Args:
+ random_state_dict: A dictionary of the form returned by `get_rng_state`.
+ """
+ random.setstate(random_state_dict["random_state"])
+ np.random.set_state(random_state_dict["numpy_random_state"])
+ torch.random.set_rng_state(random_state_dict["torch_random_state"])
+ if torch.cuda.is_available():
+ torch.cuda.random.set_rng_state(random_state_dict["torch_cuda_random_state"])
+
+
+def set_seed(seed) -> None:
+ """Set seed for reproducibility."""
+ random.seed(seed)
+ np.random.seed(seed)
+ torch.manual_seed(seed)
+ if torch.cuda.is_available():
+ torch.cuda.manual_seed_all(seed)
+
+
+@contextmanager
+def seeded_context(seed: int) -> Generator[None, None, None]:
+ """Set the seed when entering a context, and restore the prior random state at exit.
+
+ Example usage:
+
+ ```
+ a = random.random() # produces some random number
+ with seeded_context(1337):
+ b = random.random() # produces some other random number
+ c = random.random() # produces yet another random number, but the same it would have if we never made `b`
+ ```
+ """
+ random_state_dict = get_rng_state()
+ set_seed(seed)
+ yield None
+ set_rng_state(random_state_dict)
diff --git a/lerobot/common/utils/train_utils.py b/lerobot/common/utils/train_utils.py
new file mode 100644
index 00000000..a7998312
--- /dev/null
+++ b/lerobot/common/utils/train_utils.py
@@ -0,0 +1,161 @@
+#!/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 logging
+from pathlib import Path
+
+from termcolor import colored
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+
+from lerobot.common.constants import (
+ CHECKPOINTS_DIR,
+ LAST_CHECKPOINT_LINK,
+ PRETRAINED_MODEL_DIR,
+ TRAINING_STATE_DIR,
+ TRAINING_STEP,
+)
+from lerobot.common.datasets.utils import load_json, write_json
+from lerobot.common.optim.optimizers import load_optimizer_state, save_optimizer_state
+from lerobot.common.optim.schedulers import load_scheduler_state, save_scheduler_state
+from lerobot.common.policies.pretrained import PreTrainedPolicy
+from lerobot.common.utils.random_utils import load_rng_state, save_rng_state
+from lerobot.configs.train import TrainPipelineConfig
+
+
+def log_output_dir(out_dir):
+ logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {out_dir}")
+
+
+def get_step_identifier(step: int, total_steps: int) -> str:
+ num_digits = max(6, len(str(total_steps)))
+ return f"{step:0{num_digits}d}"
+
+
+def get_step_checkpoint_dir(output_dir: Path, total_steps: int, step: int) -> Path:
+ """Returns the checkpoint sub-directory corresponding to the step number."""
+ step_identifier = get_step_identifier(step, total_steps)
+ return output_dir / CHECKPOINTS_DIR / step_identifier
+
+
+def save_training_step(step: int, save_dir: Path) -> None:
+ write_json({"step": step}, save_dir / TRAINING_STEP)
+
+
+def load_training_step(save_dir: Path) -> int:
+ training_step = load_json(save_dir / TRAINING_STEP)
+ return training_step["step"]
+
+
+def update_last_checkpoint(checkpoint_dir: Path) -> Path:
+ last_checkpoint_dir = checkpoint_dir.parent / LAST_CHECKPOINT_LINK
+ if last_checkpoint_dir.is_symlink():
+ last_checkpoint_dir.unlink()
+ relative_target = checkpoint_dir.relative_to(checkpoint_dir.parent)
+ last_checkpoint_dir.symlink_to(relative_target)
+
+
+def save_checkpoint(
+ checkpoint_dir: Path,
+ step: int,
+ cfg: TrainPipelineConfig,
+ policy: PreTrainedPolicy,
+ optimizer: Optimizer,
+ scheduler: LRScheduler | None = None,
+) -> None:
+ """This function creates the following directory structure:
+
+ 005000/ # training step at checkpoint
+ ├── pretrained_model/
+ │ ├── config.json # policy config
+ │ ├── model.safetensors # policy weights
+ │ └── train_config.json # train config
+ └── training_state/
+ ├── optimizer_param_groups.json # optimizer param groups
+ ├── optimizer_state.safetensors # optimizer state
+ ├── rng_state.safetensors # rng states
+ ├── scheduler_state.json # scheduler state
+ └── training_step.json # training step
+
+ Args:
+ cfg (TrainPipelineConfig): The training config used for this run.
+ step (int): The training step at that checkpoint.
+ policy (PreTrainedPolicy): The policy to save.
+ optimizer (Optimizer | None, optional): The optimizer to save the state from. Defaults to None.
+ scheduler (LRScheduler | None, optional): The scheduler to save the state from. Defaults to None.
+ """
+ pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR
+ policy.save_pretrained(pretrained_dir)
+ cfg.save_pretrained(pretrained_dir)
+ save_training_state(checkpoint_dir, step, optimizer, scheduler)
+
+
+def save_training_state(
+ checkpoint_dir: Path,
+ train_step: int,
+ optimizer: Optimizer | None = None,
+ scheduler: LRScheduler | None = None,
+) -> None:
+ """
+ Saves the training step, optimizer state, scheduler state, and rng state.
+
+ Args:
+ save_dir (Path): The directory to save artifacts to.
+ train_step (int): Current training step.
+ optimizer (Optimizer | None, optional): The optimizer from which to save the state_dict.
+ Defaults to None.
+ scheduler (LRScheduler | None, optional): The scheduler from which to save the state_dict.
+ Defaults to None.
+ """
+ save_dir = checkpoint_dir / TRAINING_STATE_DIR
+ save_dir.mkdir(parents=True, exist_ok=True)
+ save_training_step(train_step, save_dir)
+ save_rng_state(save_dir)
+ if optimizer is not None:
+ save_optimizer_state(optimizer, save_dir)
+ if scheduler is not None:
+ save_scheduler_state(scheduler, save_dir)
+
+
+def load_training_state(
+ checkpoint_dir: Path, optimizer: Optimizer, scheduler: LRScheduler | None
+) -> tuple[int, Optimizer, LRScheduler | None]:
+ """
+ Loads the training step, optimizer state, scheduler state, and rng state.
+ This is used to resume a training run.
+
+ Args:
+ checkpoint_dir (Path): The checkpoint directory. Should contain a 'training_state' dir.
+ optimizer (Optimizer): The optimizer to load the state_dict to.
+ scheduler (LRScheduler | None): The scheduler to load the state_dict to (can be None).
+
+ Raises:
+ NotADirectoryError: If 'checkpoint_dir' doesn't contain a 'training_state' dir
+
+ Returns:
+ tuple[int, Optimizer, LRScheduler | None]: training step, optimizer and scheduler with their
+ state_dict loaded.
+ """
+ training_state_dir = checkpoint_dir / TRAINING_STATE_DIR
+ if not training_state_dir.is_dir():
+ raise NotADirectoryError(training_state_dir)
+
+ load_rng_state(training_state_dir)
+ step = load_training_step(training_state_dir)
+ optimizer = load_optimizer_state(optimizer, training_state_dir)
+ if scheduler is not None:
+ scheduler = load_scheduler_state(scheduler, training_state_dir)
+
+ return step, optimizer, scheduler
diff --git a/lerobot/common/utils/utils.py b/lerobot/common/utils/utils.py
index cb4f1874..563a7b81 100644
--- a/lerobot/common/utils/utils.py
+++ b/lerobot/common/utils/utils.py
@@ -17,12 +17,10 @@ import logging
import os
import os.path as osp
import platform
-import random
-from contextlib import contextmanager
+import subprocess
from copy import copy
from datetime import datetime, timezone
from pathlib import Path
-from typing import Any, Generator
import numpy as np
import torch
@@ -53,8 +51,10 @@ def auto_select_torch_device() -> torch.device:
return torch.device("cpu")
+# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
"""Given a string, return a torch.device with checks on whether the device is available."""
+ try_device = str(try_device)
match try_device:
case "cuda":
assert torch.cuda.is_available()
@@ -87,6 +87,7 @@ def get_safe_dtype(dtype: torch.dtype, device: str | torch.device):
def is_torch_device_available(try_device: str) -> bool:
+ try_device = str(try_device) # Ensure try_device is a string
if try_device == "cuda":
return torch.cuda.is_available()
elif try_device == "mps":
@@ -94,7 +95,7 @@ def is_torch_device_available(try_device: str) -> bool:
elif try_device == "cpu":
return True
else:
- raise ValueError(f"Unknown device '{try_device}.")
+ raise ValueError(f"Unknown device {try_device}. Supported devices are: cuda, mps or cpu.")
def is_amp_available(device: str):
@@ -106,59 +107,6 @@ def is_amp_available(device: str):
raise ValueError(f"Unknown device '{device}.")
-def get_global_random_state() -> dict[str, Any]:
- """Get the random state for `random`, `numpy`, and `torch`."""
- random_state_dict = {
- "random_state": random.getstate(),
- "numpy_random_state": np.random.get_state(),
- "torch_random_state": torch.random.get_rng_state(),
- }
- if torch.cuda.is_available():
- random_state_dict["torch_cuda_random_state"] = torch.cuda.random.get_rng_state()
- return random_state_dict
-
-
-def set_global_random_state(random_state_dict: dict[str, Any]):
- """Set the random state for `random`, `numpy`, and `torch`.
-
- Args:
- random_state_dict: A dictionary of the form returned by `get_global_random_state`.
- """
- random.setstate(random_state_dict["random_state"])
- np.random.set_state(random_state_dict["numpy_random_state"])
- torch.random.set_rng_state(random_state_dict["torch_random_state"])
- if torch.cuda.is_available():
- torch.cuda.random.set_rng_state(random_state_dict["torch_cuda_random_state"])
-
-
-def set_global_seed(seed):
- """Set seed for reproducibility."""
- random.seed(seed)
- np.random.seed(seed)
- torch.manual_seed(seed)
- if torch.cuda.is_available():
- torch.cuda.manual_seed_all(seed)
-
-
-@contextmanager
-def seeded_context(seed: int) -> Generator[None, None, None]:
- """Set the seed when entering a context, and restore the prior random state at exit.
-
- Example usage:
-
- ```
- a = random.random() # produces some random number
- with seeded_context(1337):
- b = random.random() # produces some other random number
- c = random.random() # produces yet another random number, but the same it would have if we never made `b`
- ```
- """
- random_state_dict = get_global_random_state()
- set_global_seed(seed)
- yield None
- set_global_random_state(random_state_dict)
-
-
def init_logging():
def custom_format(record):
dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
@@ -221,23 +169,31 @@ def capture_timestamp_utc():
def say(text, blocking=False):
- # Check if mac, linux, or windows.
- if platform.system() == "Darwin":
- cmd = f'say "{text}"'
- if not blocking:
- cmd += " &"
- elif platform.system() == "Linux":
- cmd = f'spd-say "{text}"'
- if blocking:
- cmd += " --wait"
- elif platform.system() == "Windows":
- # TODO(rcadene): Make blocking option work for Windows
- cmd = (
- 'PowerShell -Command "Add-Type -AssemblyName System.Speech; '
- f"(New-Object System.Speech.Synthesis.SpeechSynthesizer).Speak('{text}')\""
- )
+ system = platform.system()
- os.system(cmd)
+ if system == "Darwin":
+ cmd = ["say", text]
+
+ elif system == "Linux":
+ cmd = ["spd-say", text]
+ if blocking:
+ cmd.append("--wait")
+
+ elif system == "Windows":
+ cmd = [
+ "PowerShell",
+ "-Command",
+ "Add-Type -AssemblyName System.Speech; "
+ f"(New-Object System.Speech.Synthesis.SpeechSynthesizer).Speak('{text}')",
+ ]
+
+ else:
+ raise RuntimeError("Unsupported operating system for text-to-speech.")
+
+ if blocking:
+ subprocess.run(cmd, check=True)
+ else:
+ subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
def log_say(text, play_sounds, blocking=False):
@@ -257,5 +213,18 @@ def get_channel_first_image_shape(image_shape: tuple) -> tuple:
return shape
-def has_method(cls: object, method_name: str):
+def has_method(cls: object, method_name: str) -> bool:
return hasattr(cls, method_name) and callable(getattr(cls, method_name))
+
+
+def is_valid_numpy_dtype_string(dtype_str: str) -> bool:
+ """
+ Return True if a given string can be converted to a numpy dtype.
+ """
+ try:
+ # Attempt to convert the string to a numpy dtype
+ np.dtype(dtype_str)
+ return True
+ except TypeError:
+ # If a TypeError is raised, the string is not a valid dtype
+ return False
diff --git a/lerobot/common/utils/wandb_utils.py b/lerobot/common/utils/wandb_utils.py
new file mode 100644
index 00000000..700ebea5
--- /dev/null
+++ b/lerobot/common/utils/wandb_utils.py
@@ -0,0 +1,127 @@
+#!/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 logging
+import os
+import re
+from glob import glob
+from pathlib import Path
+
+from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
+from termcolor import colored
+
+from lerobot.common.constants import PRETRAINED_MODEL_DIR
+from lerobot.configs.train import TrainPipelineConfig
+
+
+def cfg_to_group(cfg: TrainPipelineConfig, return_list: bool = False) -> list[str] | str:
+ """Return a group name for logging. Optionally returns group name as list."""
+ lst = [
+ f"policy:{cfg.policy.type}",
+ f"dataset:{cfg.dataset.repo_id}",
+ f"seed:{cfg.seed}",
+ ]
+ if cfg.env is not None:
+ lst.append(f"env:{cfg.env.type}")
+ return lst if return_list else "-".join(lst)
+
+
+def get_wandb_run_id_from_filesystem(log_dir: Path) -> str:
+ # Get the WandB run ID.
+ paths = glob(str(log_dir / "wandb/latest-run/run-*"))
+ if len(paths) != 1:
+ raise RuntimeError("Couldn't get the previous WandB run ID for run resumption.")
+ match = re.search(r"run-([^\.]+).wandb", paths[0].split("/")[-1])
+ if match is None:
+ raise RuntimeError("Couldn't get the previous WandB run ID for run resumption.")
+ wandb_run_id = match.groups(0)[0]
+ return wandb_run_id
+
+
+def get_safe_wandb_artifact_name(name: str):
+ """WandB artifacts don't accept ":" or "/" in their name."""
+ return name.replace(":", "_").replace("/", "_")
+
+
+class WandBLogger:
+ """A helper class to log object using wandb."""
+
+ def __init__(self, cfg: TrainPipelineConfig):
+ self.cfg = cfg.wandb
+ self.log_dir = cfg.output_dir
+ self.job_name = cfg.job_name
+ self.env_fps = cfg.env.fps if cfg.env else None
+ self._group = cfg_to_group(cfg)
+
+ # Set up WandB.
+ os.environ["WANDB_SILENT"] = "True"
+ import wandb
+
+ wandb_run_id = (
+ cfg.wandb.run_id
+ if cfg.wandb.run_id
+ else get_wandb_run_id_from_filesystem(self.log_dir)
+ if cfg.resume
+ else None
+ )
+ wandb.init(
+ id=wandb_run_id,
+ project=self.cfg.project,
+ entity=self.cfg.entity,
+ name=self.job_name,
+ notes=self.cfg.notes,
+ tags=cfg_to_group(cfg, return_list=True),
+ dir=self.log_dir,
+ config=cfg.to_dict(),
+ # TODO(rcadene): try set to True
+ save_code=False,
+ # TODO(rcadene): split train and eval, and run async eval with job_type="eval"
+ job_type="train_eval",
+ resume="must" if cfg.resume else None,
+ )
+ print(colored("Logs will be synced with wandb.", "blue", attrs=["bold"]))
+ logging.info(f"Track this run --> {colored(wandb.run.get_url(), 'yellow', attrs=['bold'])}")
+ self._wandb = wandb
+
+ def log_policy(self, checkpoint_dir: Path):
+ """Checkpoints the policy to wandb."""
+ if self.cfg.disable_artifact:
+ return
+
+ step_id = checkpoint_dir.name
+ artifact_name = f"{self._group}-{step_id}"
+ artifact_name = get_safe_wandb_artifact_name(artifact_name)
+ artifact = self._wandb.Artifact(artifact_name, type="model")
+ artifact.add_file(checkpoint_dir / PRETRAINED_MODEL_DIR / SAFETENSORS_SINGLE_FILE)
+ self._wandb.log_artifact(artifact)
+
+ def log_dict(self, d: dict, step: int, mode: str = "train"):
+ if mode not in {"train", "eval"}:
+ raise ValueError(mode)
+
+ for k, v in d.items():
+ if not isinstance(v, (int, float, str)):
+ logging.warning(
+ f'WandB logging of key "{k}" was ignored as its type is not handled by this wrapper.'
+ )
+ continue
+ self._wandb.log({f"{mode}/{k}": v}, step=step)
+
+ def log_video(self, video_path: str, step: int, mode: str = "train"):
+ if mode not in {"train", "eval"}:
+ raise ValueError(mode)
+
+ wandb_video = self._wandb.Video(video_path, fps=self.env_fps, format="mp4")
+ self._wandb.log({f"{mode}/video": wandb_video}, step=step)
diff --git a/lerobot/configs/default.py b/lerobot/configs/default.py
index 5dd2f898..b23bbb6d 100644
--- a/lerobot/configs/default.py
+++ b/lerobot/configs/default.py
@@ -20,6 +20,7 @@ from lerobot.common import (
policies, # noqa: F401
)
from lerobot.common.datasets.transforms import ImageTransformsConfig
+from lerobot.common.datasets.video_utils import get_safe_default_codec
@dataclass
@@ -27,13 +28,15 @@ class DatasetConfig:
# You may provide a list of datasets here. `train.py` creates them all and concatenates them. Note: only data
# keys common between the datasets are kept. Each dataset gets and additional transform that inserts the
# "dataset_index" into the returned item. The index mapping is made according to the order in which the
- # datsets are provided.
+ # datasets are provided.
repo_id: str
+ # Root directory where the dataset will be stored (e.g. 'dataset/path').
+ root: str | None = None
episodes: list[int] | None = None
image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
- local_files_only: bool = False
+ revision: str | None = None
use_imagenet_stats: bool = True
- video_backend: str = "pyav"
+ video_backend: str = field(default_factory=get_safe_default_codec)
@dataclass
@@ -44,6 +47,7 @@ class WandBConfig:
project: str = "lerobot"
entity: str | None = None
notes: str | None = None
+ run_id: str | None = None
@dataclass
diff --git a/lerobot/configs/eval.py b/lerobot/configs/eval.py
index 11873352..16b35291 100644
--- a/lerobot/configs/eval.py
+++ b/lerobot/configs/eval.py
@@ -1,14 +1,26 @@
+# 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 datetime as dt
import logging
from dataclasses import dataclass, field
from pathlib import Path
from lerobot.common import envs, policies # noqa: F401
-from lerobot.common.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
from lerobot.configs import parser
from lerobot.configs.default import EvalConfig
from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.train import TrainPipelineConfig
@dataclass
@@ -21,11 +33,6 @@ class EvalPipelineConfig:
policy: PreTrainedConfig | None = None
output_dir: Path | None = None
job_name: str | None = None
- # TODO(rcadene, aliberts): By default, use device and use_amp values from policy checkpoint.
- device: str | None = None # cuda | cpu | mps
- # `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
- # automatic gradient scaling is used.
- use_amp: bool = False
seed: int | None = 1000
def __post_init__(self):
@@ -36,27 +43,6 @@ class EvalPipelineConfig:
self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
self.policy.pretrained_path = policy_path
- # When no device or use_amp are given, use the one from training config.
- if self.device is None or self.use_amp is None:
- train_cfg = TrainPipelineConfig.from_pretrained(policy_path)
- if self.device is None:
- self.device = train_cfg.device
- if self.use_amp is None:
- self.use_amp = train_cfg.use_amp
-
- # Automatically switch to available device if necessary
- if not is_torch_device_available(self.device):
- auto_device = auto_select_torch_device()
- logging.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
- self.device = auto_device
-
- # Automatically deactivate AMP if necessary
- if self.use_amp and not is_amp_available(self.device):
- logging.warning(
- f"Automatic Mixed Precision (amp) is not available on device '{self.device}'. Deactivating AMP."
- )
- self.use_amp = False
-
else:
logging.warning(
"No pretrained path was provided, evaluated policy will be built from scratch (random weights)."
@@ -73,11 +59,6 @@ class EvalPipelineConfig:
eval_dir = f"{now:%Y-%m-%d}/{now:%H-%M-%S}_{self.job_name}"
self.output_dir = Path("outputs/eval") / eval_dir
- if self.device is None:
- raise ValueError("Set one of the following device: cuda, cpu or mps")
- elif self.device == "cuda" and self.use_amp is None:
- raise ValueError("Set 'use_amp' to True or False.")
-
@classmethod
def __get_path_fields__(cls) -> list[str]:
"""This enables the parser to load config from the policy using `--policy.path=local/dir`"""
diff --git a/lerobot/configs/parser.py b/lerobot/configs/parser.py
index ee784877..39e31515 100644
--- a/lerobot/configs/parser.py
+++ b/lerobot/configs/parser.py
@@ -1,4 +1,19 @@
+# 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 inspect
+import pkgutil
import sys
from argparse import ArgumentError
from functools import wraps
@@ -10,6 +25,7 @@ import draccus
from lerobot.common.utils.utils import has_method
PATH_KEY = "path"
+PLUGIN_DISCOVERY_SUFFIX = "discover_packages_path"
draccus.set_config_type("json")
@@ -45,6 +61,86 @@ def parse_arg(arg_name: str, args: Sequence[str] | None = None) -> str | None:
return None
+def parse_plugin_args(plugin_arg_suffix: str, args: Sequence[str]) -> dict:
+ """Parse plugin-related arguments from command-line arguments.
+
+ This function extracts arguments from command-line arguments that match a specified suffix pattern.
+ It processes arguments in the format '--key=value' and returns them as a dictionary.
+
+ Args:
+ plugin_arg_suffix (str): The suffix to identify plugin-related arguments.
+ cli_args (Sequence[str]): A sequence of command-line arguments to parse.
+
+ Returns:
+ dict: A dictionary containing the parsed plugin arguments where:
+ - Keys are the argument names (with '--' prefix removed if present)
+ - Values are the corresponding argument values
+
+ Example:
+ >>> args = ['--env.discover_packages_path=my_package',
+ ... '--other_arg=value']
+ >>> parse_plugin_args('discover_packages_path', args)
+ {'env.discover_packages_path': 'my_package'}
+ """
+ plugin_args = {}
+ for arg in args:
+ if "=" in arg and plugin_arg_suffix in arg:
+ key, value = arg.split("=", 1)
+ # Remove leading '--' if present
+ if key.startswith("--"):
+ key = key[2:]
+ plugin_args[key] = value
+ return plugin_args
+
+
+class PluginLoadError(Exception):
+ """Raised when a plugin fails to load."""
+
+
+def load_plugin(plugin_path: str) -> None:
+ """Load and initialize a plugin from a given Python package path.
+
+ This function attempts to load a plugin by importing its package and any submodules.
+ Plugin registration is expected to happen during package initialization, i.e. when
+ the package is imported the gym environment should be registered and the config classes
+ registered with their parents using the `register_subclass` decorator.
+
+ Args:
+ plugin_path (str): The Python package path to the plugin (e.g. "mypackage.plugins.myplugin")
+
+ Raises:
+ PluginLoadError: If the plugin cannot be loaded due to import errors or if the package path is invalid.
+
+ Examples:
+ >>> load_plugin("external_plugin.core") # Loads plugin from external package
+
+ Notes:
+ - The plugin package should handle its own registration during import
+ - All submodules in the plugin package will be imported
+ - Implementation follows the plugin discovery pattern from Python packaging guidelines
+
+ See Also:
+ https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/
+ """
+ try:
+ package_module = importlib.import_module(plugin_path, __package__)
+ except (ImportError, ModuleNotFoundError) as e:
+ raise PluginLoadError(
+ f"Failed to load plugin '{plugin_path}'. Verify the path and installation: {str(e)}"
+ ) from e
+
+ def iter_namespace(ns_pkg):
+ return pkgutil.iter_modules(ns_pkg.__path__, ns_pkg.__name__ + ".")
+
+ try:
+ for _finder, pkg_name, _ispkg in iter_namespace(package_module):
+ importlib.import_module(pkg_name)
+ except ImportError as e:
+ raise PluginLoadError(
+ f"Failed to load plugin '{plugin_path}'. Verify the path and installation: {str(e)}"
+ ) from e
+
+
def get_path_arg(field_name: str, args: Sequence[str] | None = None) -> str | None:
return parse_arg(f"{field_name}.{PATH_KEY}", args)
@@ -92,10 +188,13 @@ def filter_path_args(fields_to_filter: str | list[str], args: Sequence[str] | No
def wrap(config_path: Path | None = None):
"""
- HACK: Similar to draccus.wrap but does two additional things:
+ HACK: Similar to draccus.wrap but does three additional things:
- Will remove '.path' arguments from CLI in order to process them later on.
- If a 'config_path' is passed and the main config class has a 'from_pretrained' method, will
initialize it from there to allow to fetch configs from the hub directly
+ - Will load plugins specified in the CLI arguments. These plugins will typically register
+ their own subclasses of config classes, so that draccus can find the right class to instantiate
+ from the CLI '.type' arguments
"""
def wrapper_outer(fn):
@@ -108,6 +207,14 @@ def wrap(config_path: Path | None = None):
args = args[1:]
else:
cli_args = sys.argv[1:]
+ plugin_args = parse_plugin_args(PLUGIN_DISCOVERY_SUFFIX, cli_args)
+ for plugin_cli_arg, plugin_path in plugin_args.items():
+ try:
+ load_plugin(plugin_path)
+ except PluginLoadError as e:
+ # add the relevant CLI arg to the error message
+ raise PluginLoadError(f"{e}\nFailed plugin CLI Arg: {plugin_cli_arg}") from e
+ cli_args = filter_arg(plugin_cli_arg, cli_args)
config_path_cli = parse_arg("config_path", cli_args)
if has_method(argtype, "__get_path_fields__"):
path_fields = argtype.__get_path_fields__()
diff --git a/lerobot/configs/policies.py b/lerobot/configs/policies.py
index 9b5a7c5c..022d1fb5 100644
--- a/lerobot/configs/policies.py
+++ b/lerobot/configs/policies.py
@@ -1,4 +1,18 @@
+# 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 abc
+import logging
import os
from dataclasses import dataclass, field
from pathlib import Path
@@ -12,6 +26,7 @@ from huggingface_hub.errors import HfHubHTTPError
from lerobot.common.optim.optimizers import OptimizerConfig
from lerobot.common.optim.schedulers import LRSchedulerConfig
from lerobot.common.utils.hub import HubMixin
+from lerobot.common.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
# Generic variable that is either PreTrainedConfig or a subclass thereof
@@ -40,8 +55,24 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
input_features: dict[str, PolicyFeature] = field(default_factory=dict)
output_features: dict[str, PolicyFeature] = field(default_factory=dict)
+ device: str | None = None # cuda | cpu | mp
+ # `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
+ # automatic gradient scaling is used.
+ use_amp: bool = False
+
def __post_init__(self):
self.pretrained_path = None
+ if not self.device or not is_torch_device_available(self.device):
+ auto_device = auto_select_torch_device()
+ logging.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
+ self.device = auto_device.type
+
+ # Automatically deactivate AMP if necessary
+ if self.use_amp and not is_amp_available(self.device):
+ logging.warning(
+ f"Automatic Mixed Precision (amp) is not available on device '{self.device}'. Deactivating AMP."
+ )
+ self.use_amp = False
@property
def type(self) -> str:
diff --git a/lerobot/configs/train.py b/lerobot/configs/train.py
index 3d976e81..7a787b83 100644
--- a/lerobot/configs/train.py
+++ b/lerobot/configs/train.py
@@ -1,5 +1,17 @@
+# 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 datetime as dt
-import logging
import os
from dataclasses import dataclass, field
from pathlib import Path
@@ -13,7 +25,6 @@ from lerobot.common import envs
from lerobot.common.optim import OptimizerConfig
from lerobot.common.optim.schedulers import LRSchedulerConfig
from lerobot.common.utils.hub import HubMixin
-from lerobot.common.utils.utils import auto_select_torch_device, is_amp_available
from lerobot.configs import parser
from lerobot.configs.default import DatasetConfig, EvalConfig, WandBConfig
from lerobot.configs.policies import PreTrainedConfig
@@ -21,68 +32,6 @@ from lerobot.configs.policies import PreTrainedConfig
TRAIN_CONFIG_NAME = "train_config.json"
-@dataclass
-class OfflineConfig:
- steps: int = 100_000
-
-
-@dataclass
-class OnlineConfig:
- """
- The online training loop looks something like:
-
- ```python
- for i in range(steps):
- do_online_rollout_and_update_online_buffer()
- for j in range(steps_between_rollouts):
- batch = next(dataloader_with_offline_and_online_data)
- loss = policy(batch)
- loss.backward()
- optimizer.step()
- ```
-
- Note that the online training loop adopts most of the options from the offline loop unless specified
- otherwise.
- """
-
- steps: int = 0
- # How many episodes to collect at once when we reach the online rollout part of the training loop.
- rollout_n_episodes: int = 1
- # The number of environments to use in the gym.vector.VectorEnv. This ends up also being the batch size for
- # the policy. Ideally you should set this to by an even divisor of rollout_n_episodes.
- rollout_batch_size: int = 1
- # How many optimization steps (forward, backward, optimizer step) to do between running rollouts.
- steps_between_rollouts: int | None = None
- # The proportion of online samples (vs offline samples) to include in the online training batches.
- sampling_ratio: float = 0.5
- # First seed to use for the online rollout environment. Seeds for subsequent rollouts are incremented by 1.
- env_seed: int | None = None
- # Sets the maximum number of frames that are stored in the online buffer for online training. The buffer is
- # FIFO.
- buffer_capacity: int | None = None
- # The minimum number of frames to have in the online buffer before commencing online training.
- # If buffer_seed_size > rollout_n_episodes, the rollout will be run multiple times until the
- # seed size condition is satisfied.
- buffer_seed_size: int = 0
- # Whether to run the online rollouts asynchronously. This means we can run the online training steps in
- # parallel with the rollouts. This might be advised if your GPU has the bandwidth to handle training
- # + eval + environment rendering simultaneously.
- do_rollout_async: bool = False
-
- def __post_init__(self):
- if self.steps == 0:
- return
-
- if self.steps_between_rollouts is None:
- raise ValueError(
- "'steps_between_rollouts' must be set to a positive integer, but it is currently None."
- )
- if self.env_seed is None:
- raise ValueError("'env_seed' must be set to a positive integer, but it is currently None.")
- if self.buffer_capacity is None:
- raise ValueError("'buffer_capacity' must be set to a positive integer, but it is currently None.")
-
-
@dataclass
class TrainPipelineConfig(HubMixin):
dataset: DatasetConfig
@@ -97,23 +46,18 @@ class TrainPipelineConfig(HubMixin):
# Note that when resuming a run, the default behavior is to use the configuration from the checkpoint,
# regardless of what's provided with the training command at the time of resumption.
resume: bool = False
- device: str | None = None # cuda | cpu | mp
- # `use_amp` determines whether to use Automatic Mixed Precision (AMP) for training and evaluation. With AMP,
- # automatic gradient scaling is used.
- use_amp: bool = False
# `seed` is used for training (eg: model initialization, dataset shuffling)
# AND for the evaluation environments.
seed: int | None = 1000
# Number of workers for the dataloader.
num_workers: int = 4
batch_size: int = 8
+ steps: int = 100_000
eval_freq: int = 20_000
log_freq: int = 200
save_checkpoint: bool = True
# Checkpoint is saved every `save_freq` training iterations and after the last training step.
save_freq: int = 20_000
- offline: OfflineConfig = field(default_factory=OfflineConfig)
- online: OnlineConfig = field(default_factory=OnlineConfig)
use_policy_training_preset: bool = True
optimizer: OptimizerConfig | None = None
scheduler: LRSchedulerConfig | None = None
@@ -124,18 +68,6 @@ class TrainPipelineConfig(HubMixin):
self.checkpoint_path = None
def validate(self):
- if not self.device:
- logging.warning("No device specified, trying to infer device automatically")
- device = auto_select_torch_device()
- self.device = device.type
-
- # Automatically deactivate AMP if necessary
- if self.use_amp and not is_amp_available(self.device):
- logging.warning(
- f"Automatic Mixed Precision (amp) is not available on device '{self.device}'. Deactivating AMP."
- )
- self.use_amp = False
-
# HACK: We parse again the cli args here to get the pretrained paths if there was some.
policy_path = parser.get_path_arg("policy")
if policy_path:
@@ -147,7 +79,14 @@ class TrainPipelineConfig(HubMixin):
# The entire train config is already loaded, we just need to get the checkpoint dir
config_path = parser.parse_arg("config_path")
if not config_path:
- raise ValueError("A config_path is expected when resuming a run.")
+ raise ValueError(
+ f"A config_path is expected when resuming a run. Please specify path to {TRAIN_CONFIG_NAME}"
+ )
+ if not Path(config_path).resolve().exists():
+ raise NotADirectoryError(
+ f"{config_path=} is expected to be a local path. "
+ "Resuming from the hub is not supported for now."
+ )
policy_path = Path(config_path).parent
self.policy.pretrained_path = policy_path
self.checkpoint_path = policy_path.parent
@@ -160,7 +99,7 @@ class TrainPipelineConfig(HubMixin):
if not self.resume and isinstance(self.output_dir, Path) and self.output_dir.is_dir():
raise FileExistsError(
- f"Output directory {self.output_dir} alreay exists and resume is {self.resume}. "
+ f"Output directory {self.output_dir} already exists and resume is {self.resume}. "
f"Please change your output directory so that {self.output_dir} is not overwritten."
)
elif not self.output_dir:
@@ -168,11 +107,8 @@ class TrainPipelineConfig(HubMixin):
train_dir = f"{now:%Y-%m-%d}/{now:%H-%M-%S}_{self.job_name}"
self.output_dir = Path("outputs/train") / train_dir
- if self.online.steps > 0:
- if isinstance(self.dataset.repo_id, list):
- raise NotImplementedError("Online training with LeRobotMultiDataset is not implemented.")
- if self.env is None:
- raise ValueError("An environment is required for online training")
+ if isinstance(self.dataset.repo_id, list):
+ raise NotImplementedError("LeRobotMultiDataset is not currently implemented.")
if not self.use_policy_training_preset and (self.optimizer is None or self.scheduler is None):
raise ValueError("Optimizer and Scheduler must be set when the policy presets are not used.")
@@ -185,6 +121,9 @@ class TrainPipelineConfig(HubMixin):
"""This enables the parser to load config from the policy using `--policy.path=local/dir`"""
return ["policy"]
+ def to_dict(self) -> dict:
+ return draccus.encode(self)
+
def _save_pretrained(self, save_directory: Path) -> None:
with open(save_directory / TRAIN_CONFIG_NAME, "w") as f, draccus.config_type("json"):
draccus.dump(self, f, indent=4)
diff --git a/lerobot/configs/types.py b/lerobot/configs/types.py
index 0ca45a19..6b3d92e8 100644
--- a/lerobot/configs/types.py
+++ b/lerobot/configs/types.py
@@ -1,3 +1,16 @@
+# 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.
# Note: We subclass str so that serialization is straightforward
# https://stackoverflow.com/questions/24481852/serialising-an-enum-member-to-json
from dataclasses import dataclass
diff --git a/lerobot/scripts/configure_motor.py b/lerobot/scripts/configure_motor.py
index f7e07070..b0dc8a97 100644
--- a/lerobot/scripts/configure_motor.py
+++ b/lerobot/scripts/configure_motor.py
@@ -1,3 +1,16 @@
+# 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.
"""
This script configure a single motor at a time to a given ID and baudrate.
diff --git a/lerobot/scripts/control_robot.py b/lerobot/scripts/control_robot.py
index 3fdb0acc..3c3c43f9 100644
--- a/lerobot/scripts/control_robot.py
+++ b/lerobot/scripts/control_robot.py
@@ -1,3 +1,16 @@
+# 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.
"""
Utilities to control a robot.
@@ -77,6 +90,13 @@ python lerobot/scripts/control_robot.py record \
--control.reset_time_s=10
```
+- For remote controlled robots like LeKiwi, run this script on the robot edge device (e.g. RaspBerryPi):
+```bash
+python lerobot/scripts/control_robot.py \
+ --robot.type=lekiwi \
+ --control.type=remote_robot
+```
+
**NOTE**: You can use your keyboard to control data recording flow.
- Tap right arrow key '->' to early exit while recording an episode and go to resseting the environment.
- Tap right arrow key '->' to early exit while resetting the environment and got to recording the next episode.
@@ -85,7 +105,6 @@ python lerobot/scripts/control_robot.py record \
This might require a sudo permission to allow your terminal to monitor keyboard events.
**NOTE**: You can resume/continue data recording by running the same data recording command and adding `--control.resume=true`.
-If the dataset you want to extend is not on the hub, you also need to add `--control.local_files_only=true`.
- Train on this dataset with the ACT policy:
```bash
@@ -127,6 +146,7 @@ from lerobot.common.robot_devices.control_configs import (
CalibrateControlConfig,
ControlPipelineConfig,
RecordControlConfig,
+ RemoteRobotConfig,
ReplayControlConfig,
TeleoperateControlConfig,
)
@@ -188,6 +208,16 @@ def calibrate(robot: Robot, cfg: CalibrateControlConfig):
if robot.is_connected:
robot.disconnect()
+ if robot.robot_type.startswith("lekiwi") and "main_follower" in arms:
+ print("Calibrating only the lekiwi follower arm 'main_follower'...")
+ robot.calibrate_follower()
+ return
+
+ if robot.robot_type.startswith("lekiwi") and "main_leader" in arms:
+ print("Calibrating only the lekiwi leader arm 'main_leader'...")
+ robot.calibrate_leader()
+ return
+
# Calling `connect` automatically runs calibration
# when the calibration file is missing
robot.connect()
@@ -216,7 +246,6 @@ def record(
dataset = LeRobotDataset(
cfg.repo_id,
root=cfg.root,
- local_files_only=cfg.local_files_only,
)
if len(robot.cameras) > 0:
dataset.start_image_writer(
@@ -238,7 +267,7 @@ def record(
)
# Load pretrained policy
- policy = None if cfg.policy is None else make_policy(cfg.policy, cfg.device, ds_meta=dataset.meta)
+ policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta)
if not robot.is_connected:
robot.connect()
@@ -263,15 +292,14 @@ def record(
log_say(f"Recording episode {dataset.num_episodes}", cfg.play_sounds)
record_episode(
- dataset=dataset,
robot=robot,
+ dataset=dataset,
events=events,
episode_time_s=cfg.episode_time_s,
display_cameras=cfg.display_cameras,
policy=policy,
- device=cfg.device,
- use_amp=cfg.use_amp,
fps=cfg.fps,
+ single_task=cfg.single_task,
)
# Execute a few seconds without recording to give time to manually reset the environment
@@ -282,7 +310,7 @@ def record(
(recorded_episodes < cfg.num_episodes - 1) or events["rerecord_episode"]
):
log_say("Reset the environment", cfg.play_sounds)
- reset_environment(robot, events, cfg.reset_time_s)
+ reset_environment(robot, events, cfg.reset_time_s, cfg.fps)
if events["rerecord_episode"]:
log_say("Re-record episode", cfg.play_sounds)
@@ -291,7 +319,7 @@ def record(
dataset.clear_episode_buffer()
continue
- dataset.save_episode(cfg.single_task)
+ dataset.save_episode()
recorded_episodes += 1
if events["stop_recording"]:
@@ -300,11 +328,6 @@ def record(
log_say("Stop recording", cfg.play_sounds, blocking=True)
stop_recording(robot, listener, cfg.display_cameras)
- if cfg.run_compute_stats:
- logging.info("Computing dataset statistics")
-
- dataset.consolidate(cfg.run_compute_stats)
-
if cfg.push_to_hub:
dataset.push_to_hub(tags=cfg.tags, private=cfg.private)
@@ -320,9 +343,7 @@ def replay(
# TODO(rcadene, aliberts): refactor with control_loop, once `dataset` is an instance of LeRobotDataset
# TODO(rcadene): Add option to record logs
- dataset = LeRobotDataset(
- cfg.repo_id, root=cfg.root, episodes=[cfg.episode], local_files_only=cfg.local_files_only
- )
+ dataset = LeRobotDataset(cfg.repo_id, root=cfg.root, episodes=[cfg.episode])
actions = dataset.hf_dataset.select_columns("action")
if not robot.is_connected:
@@ -357,6 +378,10 @@ def control_robot(cfg: ControlPipelineConfig):
record(robot, cfg.control)
elif isinstance(cfg.control, ReplayControlConfig):
replay(robot, cfg.control)
+ elif isinstance(cfg.control, RemoteRobotConfig):
+ from lerobot.common.robot_devices.robots.lekiwi_remote import run_lekiwi
+
+ run_lekiwi(cfg.robot)
if robot.is_connected:
# Disconnect manually to avoid a "Core dump" during process
diff --git a/lerobot/scripts/control_sim_robot.py b/lerobot/scripts/control_sim_robot.py
index 8d69bf31..5347822c 100644
--- a/lerobot/scripts/control_sim_robot.py
+++ b/lerobot/scripts/control_sim_robot.py
@@ -1,3 +1,16 @@
+# 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.
"""
Utilities to control a robot in simulation.
@@ -59,8 +72,8 @@ python lerobot/scripts/control_sim_robot.py record \
```
**NOTE**: You can use your keyboard to control data recording flow.
-- Tap right arrow key '->' to early exit while recording an episode and go to reseting the environment.
-- Tap right arrow key '->' to early exit while reseting the environment and got to recording the next episode.
+- Tap right arrow key '->' to early exit while recording an episode and go to resetting the environment.
+- Tap right arrow key '->' to early exit while resetting the environment and got to recording the next episode.
- Tap left arrow key '<-' to early exit and re-record the current episode.
- Tap escape key 'esc' to stop the data recording.
This might require a sudo permission to allow your terminal to monitor keyboard events.
@@ -131,7 +144,7 @@ def none_or_int(value):
def init_sim_calibration(robot, cfg):
# Constants necessary for transforming the joint pos of the real robot to the sim
- # depending on the robot discription used in that sim.
+ # depending on the robot description used in that sim.
start_pos = np.array(robot.leader_arms.main.calibration["start_pos"])
axis_directions = np.array(cfg.get("axis_directions", [1]))
offsets = np.array(cfg.get("offsets", [0])) * np.pi
@@ -445,7 +458,7 @@ if __name__ == "__main__":
type=int,
default=0,
help=(
- "Number of subprocesses handling the saving of frames as PNGs. Set to 0 to use threads only; "
+ "Number of subprocesses handling the saving of frames as PNG. Set to 0 to use threads only; "
"set to ≥1 to use subprocesses, each using threads to write images. The best number of processes "
"and threads depends on your system. We recommend 4 threads per camera with 0 processes. "
"If fps is unstable, adjust the thread count. If still unstable, try using 1 or more subprocesses."
diff --git a/lerobot/scripts/eval.py b/lerobot/scripts/eval.py
index 253bc45c..d7a4201f 100644
--- a/lerobot/scripts/eval.py
+++ b/lerobot/scripts/eval.py
@@ -61,21 +61,21 @@ import einops
import gymnasium as gym
import numpy as np
import torch
+from termcolor import colored
from torch import Tensor, nn
from tqdm import trange
from lerobot.common.envs.factory import make_env
from lerobot.common.envs.utils import preprocess_observation
-from lerobot.common.logger import log_output_dir
from lerobot.common.policies.factory import make_policy
from lerobot.common.policies.pretrained import PreTrainedPolicy
from lerobot.common.policies.utils import get_device_from_parameters
from lerobot.common.utils.io_utils import write_video
+from lerobot.common.utils.random_utils import set_seed
from lerobot.common.utils.utils import (
get_safe_torch_device,
init_logging,
inside_slurm,
- set_global_seed,
)
from lerobot.configs import parser
from lerobot.configs.eval import EvalPipelineConfig
@@ -125,9 +125,6 @@ def rollout(
# Reset the policy and environments.
policy.reset()
- if hasattr(policy, "use_ema_modules"):
- policy.use_ema_modules()
-
observation, info = env.reset(seed=seeds)
if render_callback is not None:
render_callback(env)
@@ -154,7 +151,9 @@ def rollout(
if return_observations:
all_observations.append(deepcopy(observation))
- observation = {key: observation[key].to(device, non_blocking=True) for key in observation}
+ observation = {
+ key: observation[key].to(device, non_blocking=device.type == "cuda") for key in observation
+ }
with torch.inference_mode():
action = policy.select_action(observation)
@@ -455,30 +454,30 @@ def _compile_episode_data(
@parser.wrap()
-def eval(cfg: EvalPipelineConfig):
+def eval_main(cfg: EvalPipelineConfig):
logging.info(pformat(asdict(cfg)))
# Check device is available
- device = get_safe_torch_device(cfg.device, log=True)
+ device = get_safe_torch_device(cfg.policy.device, log=True)
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
- set_global_seed(cfg.seed)
+ set_seed(cfg.seed)
- log_output_dir(cfg.output_dir)
+ logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
logging.info("Making environment.")
env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
logging.info("Making policy.")
+
policy = make_policy(
cfg=cfg.policy,
- device=device,
env_cfg=cfg.env,
)
policy.eval()
- with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.use_amp else nullcontext():
+ with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
info = eval_policy(
env,
policy,
@@ -500,4 +499,4 @@ def eval(cfg: EvalPipelineConfig):
if __name__ == "__main__":
init_logging()
- eval()
+ eval_main()
diff --git a/lerobot/scripts/find_motors_bus_port.py b/lerobot/scripts/find_motors_bus_port.py
index 67b92ad7..68f2315d 100644
--- a/lerobot/scripts/find_motors_bus_port.py
+++ b/lerobot/scripts/find_motors_bus_port.py
@@ -1,3 +1,16 @@
+# 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 os
import time
from pathlib import Path
diff --git a/lerobot/scripts/push_dataset_to_hub.py b/lerobot/scripts/push_dataset_to_hub.py
index 0233ede6..3de2462b 100644
--- a/lerobot/scripts/push_dataset_to_hub.py
+++ b/lerobot/scripts/push_dataset_to_hub.py
@@ -175,7 +175,7 @@ def push_dataset_to_hub(
# Robustify when `local_dir` is str instead of Path
local_dir = Path(local_dir)
- # Send warning if local_dir isn't well formated
+ # Send warning if local_dir isn't well formatted
if local_dir.parts[-2] != user_id or local_dir.parts[-1] != dataset_id:
warnings.warn(
f"`local_dir` ({local_dir}) doesn't contain a community or user id `/` the name of the dataset that match the `repo_id` (e.g. 'data/lerobot/pusht'). Following this naming convention is advised, but not mandatory.",
diff --git a/lerobot/scripts/train.py b/lerobot/scripts/train.py
index 9af1a972..f2b1e29e 100644
--- a/lerobot/scripts/train.py
+++ b/lerobot/scripts/train.py
@@ -15,61 +15,64 @@
# limitations under the License.
import logging
import time
-from concurrent.futures import ThreadPoolExecutor
from contextlib import nullcontext
-from copy import deepcopy
-from dataclasses import asdict
from pprint import pformat
-from threading import Lock
+from typing import Any
-import numpy as np
import torch
+from termcolor import colored
from torch.amp import GradScaler
+from torch.optim import Optimizer
-from lerobot.common.datasets.factory import make_dataset, resolve_delta_timestamps
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.datasets.online_buffer import OnlineBuffer, compute_sampler_weights
+from lerobot.common.datasets.factory import make_dataset
from lerobot.common.datasets.sampler import EpisodeAwareSampler
from lerobot.common.datasets.utils import cycle
from lerobot.common.envs.factory import make_env
-from lerobot.common.logger import Logger, log_output_dir
-from lerobot.common.optim.factory import load_training_state, make_optimizer_and_scheduler
+from lerobot.common.optim.factory import make_optimizer_and_scheduler
from lerobot.common.policies.factory import make_policy
+from lerobot.common.policies.pretrained import PreTrainedPolicy
from lerobot.common.policies.utils import get_device_from_parameters
+from lerobot.common.utils.logging_utils import AverageMeter, MetricsTracker
+from lerobot.common.utils.random_utils import set_seed
+from lerobot.common.utils.train_utils import (
+ get_step_checkpoint_dir,
+ get_step_identifier,
+ load_training_state,
+ save_checkpoint,
+ update_last_checkpoint,
+)
from lerobot.common.utils.utils import (
format_big_number,
- get_safe_dtype,
get_safe_torch_device,
has_method,
init_logging,
- set_global_seed,
)
+from lerobot.common.utils.wandb_utils import WandBLogger
from lerobot.configs import parser
from lerobot.configs.train import TrainPipelineConfig
from lerobot.scripts.eval import eval_policy
def update_policy(
- policy,
- batch,
- optimizer,
- grad_clip_norm,
+ train_metrics: MetricsTracker,
+ policy: PreTrainedPolicy,
+ batch: Any,
+ optimizer: Optimizer,
+ grad_clip_norm: float,
grad_scaler: GradScaler,
lr_scheduler=None,
use_amp: bool = False,
lock=None,
-):
- """Returns a dictionary of items for logging."""
+) -> tuple[MetricsTracker, dict]:
start_time = time.perf_counter()
device = get_device_from_parameters(policy)
policy.train()
with torch.autocast(device_type=device.type) if use_amp else nullcontext():
- output_dict = policy.forward(batch)
+ loss, output_dict = policy.forward(batch)
# TODO(rcadene): policy.unnormalize_outputs(out_dict)
- loss = output_dict["loss"]
grad_scaler.scale(loss).backward()
- # Unscale the graident of the optimzer's assigned params in-place **prior to gradient clipping**.
+ # Unscale the gradient of the optimizer's assigned params in-place **prior to gradient clipping**.
grad_scaler.unscale_(optimizer)
grad_norm = torch.nn.utils.clip_grad_norm_(
@@ -87,9 +90,6 @@ def update_policy(
optimizer.zero_grad()
- if hasattr(policy, "update_ema_modules"):
- policy.update_ema_modules()
-
# Step through pytorch scheduler at every batch instead of epoch
if lr_scheduler is not None:
lr_scheduler.step()
@@ -98,113 +98,34 @@ def update_policy(
# To possibly update an internal buffer (for instance an Exponential Moving Average like in TDMPC).
policy.update()
- info = {
- "loss": loss.item(),
- "grad_norm": float(grad_norm),
- "lr": optimizer.param_groups[0]["lr"],
- "update_s": time.perf_counter() - start_time,
- **{k: v for k, v in output_dict.items() if k != "loss"},
- }
- info.update({k: v for k, v in output_dict.items() if k not in info})
-
- return info
-
-
-def log_train_info(
- logger: Logger, info: dict, step: int, cfg: TrainPipelineConfig, dataset: LeRobotDataset, is_online: bool
-):
- loss = info["loss"]
- grad_norm = info["grad_norm"]
- lr = info["lr"]
- update_s = info["update_s"]
- dataloading_s = info["dataloading_s"]
-
- # A sample is an (observation,action) pair, where observation and action
- # can be on multiple timestamps. In a batch, we have `batch_size`` number of samples.
- num_samples = (step + 1) * cfg.batch_size
- avg_samples_per_ep = dataset.num_frames / dataset.num_episodes
- num_episodes = num_samples / avg_samples_per_ep
- num_epochs = num_samples / dataset.num_frames
- log_items = [
- f"step:{format_big_number(step)}",
- # number of samples seen during training
- f"smpl:{format_big_number(num_samples)}",
- # number of episodes seen during training
- f"ep:{format_big_number(num_episodes)}",
- # number of time all unique samples are seen
- f"epch:{num_epochs:.2f}",
- f"loss:{loss:.3f}",
- f"grdn:{grad_norm:.3f}",
- f"lr:{lr:0.1e}",
- # in seconds
- f"updt_s:{update_s:.3f}",
- f"data_s:{dataloading_s:.3f}", # if not ~0, you are bottlenecked by cpu or io
- ]
- logging.info(" ".join(log_items))
-
- info["step"] = step
- info["num_samples"] = num_samples
- info["num_episodes"] = num_episodes
- info["num_epochs"] = num_epochs
- info["is_online"] = is_online
-
- logger.log_dict(info, step, mode="train")
-
-
-def log_eval_info(logger, info, step, cfg, dataset, is_online):
- eval_s = info["eval_s"]
- avg_sum_reward = info["avg_sum_reward"]
- pc_success = info["pc_success"]
-
- # A sample is an (observation,action) pair, where observation and action
- # can be on multiple timestamps. In a batch, we have `batch_size`` number of samples.
- num_samples = (step + 1) * cfg.batch_size
- avg_samples_per_ep = dataset.num_frames / dataset.num_episodes
- num_episodes = num_samples / avg_samples_per_ep
- num_epochs = num_samples / dataset.num_frames
- log_items = [
- f"step:{format_big_number(step)}",
- # number of samples seen during training
- f"smpl:{format_big_number(num_samples)}",
- # number of episodes seen during training
- f"ep:{format_big_number(num_episodes)}",
- # number of time all unique samples are seen
- f"epch:{num_epochs:.2f}",
- f"∑rwrd:{avg_sum_reward:.3f}",
- f"success:{pc_success:.1f}%",
- f"eval_s:{eval_s:.3f}",
- ]
- logging.info(" ".join(log_items))
-
- info["step"] = step
- info["num_samples"] = num_samples
- info["num_episodes"] = num_episodes
- info["num_epochs"] = num_epochs
- info["is_online"] = is_online
-
- logger.log_dict(info, step, mode="eval")
+ train_metrics.loss = loss.item()
+ train_metrics.grad_norm = grad_norm.item()
+ train_metrics.lr = optimizer.param_groups[0]["lr"]
+ train_metrics.update_s = time.perf_counter() - start_time
+ return train_metrics, output_dict
@parser.wrap()
def train(cfg: TrainPipelineConfig):
cfg.validate()
+ logging.info(pformat(cfg.to_dict()))
- logging.info(pformat(asdict(cfg)))
-
- # log metrics to terminal and wandb
- logger = Logger(cfg)
+ if cfg.wandb.enable and cfg.wandb.project:
+ wandb_logger = WandBLogger(cfg)
+ else:
+ wandb_logger = None
+ logging.info(colored("Logs will be saved locally.", "yellow", attrs=["bold"]))
if cfg.seed is not None:
- set_global_seed(cfg.seed)
+ set_seed(cfg.seed)
# Check device is available
- device = get_safe_torch_device(cfg.device, log=True)
-
+ device = get_safe_torch_device(cfg.policy.device, log=True)
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
logging.info("Creating dataset")
- offline_dataset = make_dataset(cfg)
+ dataset = make_dataset(cfg)
# Create environment used for evaluating checkpoints during training on simulation data.
# On real-world data, no need to create an environment as evaluations are done outside train.py,
@@ -217,13 +138,12 @@ def train(cfg: TrainPipelineConfig):
logging.info("Creating policy")
policy = make_policy(
cfg=cfg.policy,
- device=device,
- ds_meta=offline_dataset.meta,
+ ds_meta=dataset.meta,
)
logging.info("Creating optimizer and scheduler")
optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
- grad_scaler = GradScaler(device, enabled=cfg.use_amp)
+ grad_scaler = GradScaler(device.type, enabled=cfg.policy.use_amp)
step = 0 # number of policy updates (forward + backward + optim)
@@ -233,65 +153,29 @@ def train(cfg: TrainPipelineConfig):
num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad)
num_total_params = sum(p.numel() for p in policy.parameters())
- log_output_dir(cfg.output_dir)
+ logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
if cfg.env is not None:
logging.info(f"{cfg.env.task=}")
- logging.info(f"{cfg.offline.steps=} ({format_big_number(cfg.offline.steps)})")
- logging.info(f"{cfg.online.steps=}")
- logging.info(f"{offline_dataset.num_frames=} ({format_big_number(offline_dataset.num_frames)})")
- logging.info(f"{offline_dataset.num_episodes=}")
+ logging.info(f"{cfg.steps=} ({format_big_number(cfg.steps)})")
+ logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})")
+ logging.info(f"{dataset.num_episodes=}")
logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})")
logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
- # Note: this helper will be used in offline and online training loops.
- def evaluate_and_checkpoint_if_needed(step: int, is_online: bool):
- _num_digits = max(6, len(str(cfg.offline.steps + cfg.online.steps)))
- step_identifier = f"{step:0{_num_digits}d}"
-
- if cfg.env is not None and cfg.eval_freq > 0 and step % cfg.eval_freq == 0:
- logging.info(f"Eval policy at step {step}")
- with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.use_amp else nullcontext():
- eval_info = eval_policy(
- eval_env,
- policy,
- cfg.eval.n_episodes,
- videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_identifier}",
- max_episodes_rendered=4,
- start_seed=cfg.seed,
- )
- log_eval_info(logger, eval_info["aggregated"], step, cfg, offline_dataset, is_online=is_online)
- if cfg.wandb.enable:
- logger.log_video(eval_info["video_paths"][0], step, mode="eval")
- logging.info("Resume training")
-
- if cfg.save_checkpoint and (
- step % cfg.save_freq == 0 or step == cfg.offline.steps + cfg.online.steps
- ):
- logging.info(f"Checkpoint policy after step {step}")
- # Note: Save with step as the identifier, and format it to have at least 6 digits but more if
- # needed (choose 6 as a minimum for consistency without being overkill).
- logger.save_checkpoint(
- step,
- step_identifier,
- policy,
- optimizer,
- lr_scheduler,
- )
- logging.info("Resume training")
-
# create dataloader for offline training
- if getattr(cfg.policy, "drop_n_last_frames", None):
+ if hasattr(cfg.policy, "drop_n_last_frames"):
shuffle = False
sampler = EpisodeAwareSampler(
- offline_dataset.episode_data_index,
+ dataset.episode_data_index,
drop_n_last_frames=cfg.policy.drop_n_last_frames,
shuffle=True,
)
else:
shuffle = True
sampler = None
+
dataloader = torch.utils.data.DataLoader(
- offline_dataset,
+ dataset,
num_workers=cfg.num_workers,
batch_size=cfg.batch_size,
shuffle=shuffle,
@@ -303,257 +187,96 @@ def train(cfg: TrainPipelineConfig):
policy.train()
- if hasattr(policy, "init_ema_modules"):
- policy.init_ema_modules()
+ train_metrics = {
+ "loss": AverageMeter("loss", ":.3f"),
+ "grad_norm": AverageMeter("grdn", ":.3f"),
+ "lr": AverageMeter("lr", ":0.1e"),
+ "update_s": AverageMeter("updt_s", ":.3f"),
+ "dataloading_s": AverageMeter("data_s", ":.3f"),
+ }
- offline_step = 0
- for _ in range(step, cfg.offline.steps):
- if offline_step == 0:
- logging.info("Start offline training on a fixed dataset")
+ train_tracker = MetricsTracker(
+ cfg.batch_size, dataset.num_frames, dataset.num_episodes, train_metrics, initial_step=step
+ )
+ logging.info("Start offline training on a fixed dataset")
+ for _ in range(step, cfg.steps):
start_time = time.perf_counter()
batch = next(dl_iter)
- dataloading_s = time.perf_counter() - start_time
+ train_tracker.dataloading_s = time.perf_counter() - start_time
for key in batch:
if isinstance(batch[key], torch.Tensor):
batch[key] = batch[key].to(device, non_blocking=True)
- train_info = update_policy(
+ train_tracker, output_dict = update_policy(
+ train_tracker,
policy,
batch,
optimizer,
cfg.optimizer.grad_clip_norm,
grad_scaler=grad_scaler,
lr_scheduler=lr_scheduler,
- use_amp=cfg.use_amp,
+ use_amp=cfg.policy.use_amp,
)
- train_info["dataloading_s"] = dataloading_s
-
- if step % cfg.log_freq == 0:
- log_train_info(logger, train_info, step, cfg, offline_dataset, is_online=False)
-
- # Note: evaluate_and_checkpoint_if_needed happens **after** the `step`th training update has completed,
- # so we pass in step + 1.
- evaluate_and_checkpoint_if_needed(step + 1, is_online=False)
-
+ # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
+ # increment `step` here.
step += 1
- offline_step += 1 # noqa: SIM113
+ train_tracker.step()
+ is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0
+ is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps
+ is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0
- if cfg.online.steps == 0:
- if eval_env:
- eval_env.close()
- logging.info("End of training")
- return
+ if is_log_step:
+ logging.info(train_tracker)
+ if wandb_logger:
+ wandb_log_dict = train_tracker.to_dict()
+ if output_dict:
+ wandb_log_dict.update(output_dict)
+ wandb_logger.log_dict(wandb_log_dict, step)
+ train_tracker.reset_averages()
- # Online training.
+ if cfg.save_checkpoint and is_saving_step:
+ logging.info(f"Checkpoint policy after step {step}")
+ checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step)
+ save_checkpoint(checkpoint_dir, step, cfg, policy, optimizer, lr_scheduler)
+ update_last_checkpoint(checkpoint_dir)
+ if wandb_logger:
+ wandb_logger.log_policy(checkpoint_dir)
- # Create an env dedicated to online episodes collection from policy rollout.
- online_env = make_env(cfg.env, n_envs=cfg.online.rollout_batch_size)
- delta_timestamps = resolve_delta_timestamps(cfg.policy, offline_dataset.meta)
- online_buffer_path = logger.log_dir / "online_buffer"
- if cfg.resume and not online_buffer_path.exists():
- # If we are resuming a run, we default to the data shapes and buffer capacity from the saved online
- # buffer.
- logging.warning(
- "When online training is resumed, we load the latest online buffer from the prior run, "
- "and this might not coincide with the state of the buffer as it was at the moment the checkpoint "
- "was made. This is because the online buffer is updated on disk during training, independently "
- "of our explicit checkpointing mechanisms."
- )
- online_dataset = OnlineBuffer(
- online_buffer_path,
- data_spec={
- **{
- key: {"shape": ft.shape, "dtype": np.dtype("float32")}
- for key, ft in policy.config.input_features.items()
- },
- **{
- key: {"shape": ft.shape, "dtype": np.dtype("float32")}
- for key, ft in policy.config.output_features.items()
- },
- "next.reward": {"shape": (), "dtype": np.dtype("float32")},
- "next.done": {"shape": (), "dtype": np.dtype("?")},
- "task_index": {"shape": (), "dtype": np.dtype("int64")},
- # FIXME: 'task' is a string
- # "task": {"shape": (), "dtype": np.dtype("?")},
- # FIXME: 'next.success' is expected by pusht env but not xarm
- "next.success": {"shape": (), "dtype": np.dtype("?")},
- },
- buffer_capacity=cfg.online.buffer_capacity,
- fps=online_env.unwrapped.metadata["render_fps"],
- delta_timestamps=delta_timestamps,
- )
-
- # If we are doing online rollouts asynchronously, deepcopy the policy to use for online rollouts (this
- # makes it possible to do online rollouts in parallel with training updates).
- online_rollout_policy = deepcopy(policy) if cfg.online.do_rollout_async else policy
-
- # Create dataloader for online training.
- concat_dataset = torch.utils.data.ConcatDataset([offline_dataset, online_dataset])
- sampler_weights = compute_sampler_weights(
- offline_dataset,
- offline_drop_n_last_frames=getattr(cfg.policy, "drop_n_last_frames", 0),
- online_dataset=online_dataset,
- # +1 because online rollouts return an extra frame for the "final observation". Note: we don't have
- # this final observation in the offline datasets, but we might add them in future.
- online_drop_n_last_frames=getattr(cfg.policy, "drop_n_last_frames", 0) + 1,
- online_sampling_ratio=cfg.online.sampling_ratio,
- )
- sampler = torch.utils.data.WeightedRandomSampler(
- sampler_weights,
- num_samples=len(concat_dataset),
- replacement=True,
- )
- dataloader = torch.utils.data.DataLoader(
- concat_dataset,
- batch_size=cfg.batch_size,
- num_workers=cfg.num_workers,
- sampler=sampler,
- pin_memory=device.type != "cpu",
- drop_last=True,
- )
- dl_iter = cycle(dataloader)
-
- if cfg.online.do_rollout_async:
- # Lock and thread pool executor for asynchronous online rollouts.
- lock = Lock()
- # Note: 1 worker because we only ever want to run one set of online rollouts at a time. Batch
- # parallelization of rollouts is handled within the job.
- executor = ThreadPoolExecutor(max_workers=1)
- else:
- lock = None
-
- online_step = 0
- online_rollout_s = 0 # time take to do online rollout
- update_online_buffer_s = 0 # time taken to update the online buffer with the online rollout data
- # Time taken waiting for the online buffer to finish being updated. This is relevant when using the async
- # online rollout option.
- await_update_online_buffer_s = 0
- rollout_start_seed = cfg.online.env_seed
-
- while True:
- if online_step == cfg.online.steps:
- break
-
- if online_step == 0:
- logging.info("Start online training by interacting with environment")
-
- def sample_trajectory_and_update_buffer():
- nonlocal rollout_start_seed
-
- with lock if lock is not None else nullcontext():
- online_rollout_policy.load_state_dict(policy.state_dict())
-
- online_rollout_policy.eval()
- start_rollout_time = time.perf_counter()
-
- with torch.no_grad():
+ if cfg.env and is_eval_step:
+ step_id = get_step_identifier(step, cfg.steps)
+ logging.info(f"Eval policy at step {step}")
+ with (
+ torch.no_grad(),
+ torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(),
+ ):
eval_info = eval_policy(
- online_env,
- online_rollout_policy,
- n_episodes=cfg.online.rollout_n_episodes,
- max_episodes_rendered=min(10, cfg.online.rollout_n_episodes),
- videos_dir=logger.log_dir / "online_rollout_videos",
- return_episode_data=True,
- start_seed=(rollout_start_seed := (rollout_start_seed + cfg.batch_size) % 1000000),
+ eval_env,
+ policy,
+ cfg.eval.n_episodes,
+ videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}",
+ max_episodes_rendered=4,
+ start_seed=cfg.seed,
)
- online_rollout_s = time.perf_counter() - start_rollout_time
- if len(offline_dataset.meta.tasks) > 1:
- raise NotImplementedError("Add support for multi task.")
-
- # TODO(rcadene, aliberts): Hack to add a task to the online_dataset (0 is the first task of the offline_dataset)
- total_num_frames = eval_info["episodes"]["index"].shape[0]
- eval_info["episodes"]["task_index"] = torch.tensor([0] * total_num_frames, dtype=torch.int64)
- eval_info["episodes"]["task"] = ["do the thing"] * total_num_frames
-
- with lock if lock is not None else nullcontext():
- start_update_buffer_time = time.perf_counter()
- online_dataset.add_data(eval_info["episodes"])
-
- # Update the concatenated dataset length used during sampling.
- concat_dataset.cumulative_sizes = concat_dataset.cumsum(concat_dataset.datasets)
-
- # Update the sampling weights.
- sampler.weights = compute_sampler_weights(
- offline_dataset,
- offline_drop_n_last_frames=getattr(cfg.policy, "drop_n_last_frames", 0),
- online_dataset=online_dataset,
- # +1 because online rollouts return an extra frame for the "final observation". Note: we don't have
- # this final observation in the offline datasets, but we might add them in future.
- online_drop_n_last_frames=getattr(cfg.policy, "drop_n_last_frames", 0) + 1,
- online_sampling_ratio=cfg.online.sampling_ratio,
- )
- sampler.num_frames = len(concat_dataset)
-
- update_online_buffer_s = time.perf_counter() - start_update_buffer_time
-
- return online_rollout_s, update_online_buffer_s
-
- if lock is None:
- online_rollout_s, update_online_buffer_s = sample_trajectory_and_update_buffer()
- else:
- future = executor.submit(sample_trajectory_and_update_buffer)
- # If we aren't doing async rollouts, or if we haven't yet gotten enough examples in our buffer, wait
- # here until the rollout and buffer update is done, before proceeding to the policy update steps.
- if len(online_dataset) <= cfg.online.buffer_seed_size:
- online_rollout_s, update_online_buffer_s = future.result()
-
- if len(online_dataset) <= cfg.online.buffer_seed_size:
- logging.info(f"Seeding online buffer: {len(online_dataset)}/{cfg.online.buffer_seed_size}")
- continue
-
- policy.train()
- for _ in range(cfg.online.steps_between_rollouts):
- with lock if lock is not None else nullcontext():
- start_time = time.perf_counter()
- batch = next(dl_iter)
- dataloading_s = time.perf_counter() - start_time
-
- for key in batch:
- if isinstance(batch[key], torch.Tensor):
- dtype = get_safe_dtype(batch[key].dtype, device)
- batch[key] = batch[key].to(device=device, dtype=dtype, non_blocking=True)
-
- train_info = update_policy(
- policy,
- batch,
- optimizer,
- cfg.optimizer.grad_clip_norm,
- grad_scaler=grad_scaler,
- lr_scheduler=lr_scheduler,
- use_amp=cfg.use_amp,
- lock=lock,
+ eval_metrics = {
+ "avg_sum_reward": AverageMeter("∑rwrd", ":.3f"),
+ "pc_success": AverageMeter("success", ":.1f"),
+ "eval_s": AverageMeter("eval_s", ":.3f"),
+ }
+ eval_tracker = MetricsTracker(
+ cfg.batch_size, dataset.num_frames, dataset.num_episodes, eval_metrics, initial_step=step
)
-
- train_info["dataloading_s"] = dataloading_s
- train_info["online_rollout_s"] = online_rollout_s
- train_info["update_online_buffer_s"] = update_online_buffer_s
- train_info["await_update_online_buffer_s"] = await_update_online_buffer_s
- with lock if lock is not None else nullcontext():
- train_info["online_buffer_size"] = len(online_dataset)
-
- if step % cfg.log_freq == 0:
- log_train_info(logger, train_info, step, cfg, online_dataset, is_online=True)
-
- # Note: evaluate_and_checkpoint_if_needed happens **after** the `step`th training update has completed,
- # so we pass in step + 1.
- evaluate_and_checkpoint_if_needed(step + 1, is_online=True)
-
- step += 1
- online_step += 1
-
- # If we're doing async rollouts, we should now wait until we've completed them before proceeding
- # to do the next batch of rollouts.
- if cfg.online.do_rollout_async and future.running():
- start = time.perf_counter()
- online_rollout_s, update_online_buffer_s = future.result()
- await_update_online_buffer_s = time.perf_counter() - start
-
- if online_step >= cfg.online.steps:
- break
+ eval_tracker.eval_s = eval_info["aggregated"].pop("eval_s")
+ eval_tracker.avg_sum_reward = eval_info["aggregated"].pop("avg_sum_reward")
+ eval_tracker.pc_success = eval_info["aggregated"].pop("pc_success")
+ logging.info(eval_tracker)
+ if wandb_logger:
+ wandb_log_dict = {**eval_tracker.to_dict(), **eval_info}
+ wandb_logger.log_dict(wandb_log_dict, step, mode="eval")
+ wandb_logger.log_video(eval_info["video_paths"][0], step, mode="eval")
if eval_env:
eval_env.close()
diff --git a/lerobot/scripts/visualize_dataset.py b/lerobot/scripts/visualize_dataset.py
index ca176407..cdfea6b8 100644
--- a/lerobot/scripts/visualize_dataset.py
+++ b/lerobot/scripts/visualize_dataset.py
@@ -111,9 +111,9 @@ def visualize_dataset(
output_dir: Path | None = None,
) -> Path | None:
if save:
- assert (
- output_dir is not None
- ), "Set an output directory where to write .rrd files with `--output-dir path/to/directory`."
+ assert output_dir is not None, (
+ "Set an output directory where to write .rrd files with `--output-dir path/to/directory`."
+ )
repo_id = dataset.repo_id
@@ -207,12 +207,6 @@ def main():
required=True,
help="Episode to visualize.",
)
- parser.add_argument(
- "--local-files-only",
- type=int,
- default=0,
- help="Use local files only. By default, this script will try to fetch the dataset from the hub if it exists.",
- )
parser.add_argument(
"--root",
type=Path,
@@ -271,14 +265,25 @@ def main():
),
)
+ parser.add_argument(
+ "--tolerance-s",
+ type=float,
+ default=1e-4,
+ help=(
+ "Tolerance in seconds used to ensure data timestamps respect the dataset fps value"
+ "This is argument passed to the constructor of LeRobotDataset and maps to its tolerance_s constructor argument"
+ "If not given, defaults to 1e-4."
+ ),
+ )
+
args = parser.parse_args()
kwargs = vars(args)
repo_id = kwargs.pop("repo_id")
root = kwargs.pop("root")
- local_files_only = kwargs.pop("local_files_only")
+ tolerance_s = kwargs.pop("tolerance_s")
logging.info("Loading dataset")
- dataset = LeRobotDataset(repo_id, root=root, local_files_only=local_files_only)
+ dataset = LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)
visualize_dataset(dataset, **vars(args))
diff --git a/lerobot/scripts/visualize_dataset_html.py b/lerobot/scripts/visualize_dataset_html.py
index bfecbcdd..a89701e4 100644
--- a/lerobot/scripts/visualize_dataset_html.py
+++ b/lerobot/scripts/visualize_dataset_html.py
@@ -150,7 +150,7 @@ def run_server(
400,
)
dataset_version = (
- dataset.meta._version if isinstance(dataset, LeRobotDataset) else dataset.codebase_version
+ str(dataset.meta._version) if isinstance(dataset, LeRobotDataset) else dataset.codebase_version
)
match = re.search(r"v(\d+)\.", dataset_version)
if match:
@@ -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
@@ -194,7 +194,7 @@ def run_server(
]
response = requests.get(
- f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/episodes.jsonl"
+ f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/episodes.jsonl", timeout=5
)
response.raise_for_status()
# Split into lines and parse each line as JSON
@@ -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)
@@ -233,9 +234,17 @@ def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index)
This file will be loaded by Dygraph javascript to plot data in real time."""
columns = []
- selected_columns = [col for col, ft in dataset.features.items() if ft["dtype"] == "float32"]
+ selected_columns = [col for col, ft in dataset.features.items() if ft["dtype"] in ["float32", "int32"]]
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"]
@@ -245,16 +254,17 @@ def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index)
if isinstance(dataset, LeRobotDataset)
else dataset.features[column_name].shape[0]
)
- header += [f"{column_name}_{i}" for i in range(dim_state)]
if "names" in dataset.features[column_name] and dataset.features[column_name]["names"]:
column_names = dataset.features[column_name]["names"]
while not isinstance(column_names, list):
column_names = list(column_names.values())[0]
else:
- column_names = [f"motor_{i}" for i in range(dim_state)]
+ column_names = [f"{column_name}_{i}" for i in range(dim_state)]
columns.append({"key": column_name, "value": column_names})
+ header += column_names
+
selected_columns.insert(0, "timestamp")
if isinstance(dataset, LeRobotDataset):
@@ -290,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]:
@@ -317,7 +327,9 @@ def get_episode_language_instruction(dataset: LeRobotDataset, ep_index: int) ->
def get_dataset_info(repo_id: str) -> IterableNamespace:
- response = requests.get(f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/info.json")
+ response = requests.get(
+ f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/info.json", timeout=5
+ )
response.raise_for_status() # Raises an HTTPError for bad responses
dataset_info = response.json()
dataset_info["repo_id"] = repo_id
@@ -364,7 +376,7 @@ def visualize_dataset_html(
template_folder=template_dir,
)
else:
- # Create a simlink from the dataset video folder containg mp4 files to the output directory
+ # Create a simlink from the dataset video folder containing mp4 files to the output directory
# so that the http server can get access to the mp4 files.
if isinstance(dataset, LeRobotDataset):
ln_videos_dir = static_dir / "videos"
@@ -384,12 +396,6 @@ def main():
default=None,
help="Name of hugging face repositery containing a LeRobotDataset dataset (e.g. `lerobot/pusht` for https://huggingface.co/datasets/lerobot/pusht).",
)
- parser.add_argument(
- "--local-files-only",
- type=int,
- default=0,
- help="Use local files only. By default, this script will try to fetch the dataset from the hub if it exists.",
- )
parser.add_argument(
"--root",
type=Path,
@@ -440,17 +446,28 @@ def main():
help="Delete the output directory if it exists already.",
)
+ parser.add_argument(
+ "--tolerance-s",
+ type=float,
+ default=1e-4,
+ help=(
+ "Tolerance in seconds used to ensure data timestamps respect the dataset fps value"
+ "This is argument passed to the constructor of LeRobotDataset and maps to its tolerance_s constructor argument"
+ "If not given, defaults to 1e-4."
+ ),
+ )
+
args = parser.parse_args()
kwargs = vars(args)
repo_id = kwargs.pop("repo_id")
load_from_hf_hub = kwargs.pop("load_from_hf_hub")
root = kwargs.pop("root")
- local_files_only = kwargs.pop("local_files_only")
+ tolerance_s = kwargs.pop("tolerance_s")
dataset = None
if repo_id:
dataset = (
- LeRobotDataset(repo_id, root=root, local_files_only=local_files_only)
+ LeRobotDataset(repo_id, root=root, tolerance_s=tolerance_s)
if not load_from_hf_hub
else get_dataset_info(repo_id)
)
diff --git a/lerobot/scripts/visualize_image_transforms.py b/lerobot/scripts/visualize_image_transforms.py
index 727fe178..80935d32 100644
--- a/lerobot/scripts/visualize_image_transforms.py
+++ b/lerobot/scripts/visualize_image_transforms.py
@@ -109,7 +109,7 @@ def visualize_image_transforms(cfg: DatasetConfig, output_dir: Path = OUTPUT_DIR
dataset = LeRobotDataset(
repo_id=cfg.repo_id,
episodes=cfg.episodes,
- local_files_only=cfg.local_files_only,
+ revision=cfg.revision,
video_backend=cfg.video_backend,
)
diff --git a/lerobot/templates/visualize_dataset_template.html b/lerobot/templates/visualize_dataset_template.html
index 08de3e3d..cf9d40f1 100644
--- a/lerobot/templates/visualize_dataset_template.html
+++ b/lerobot/templates/visualize_dataset_template.html
@@ -14,21 +14,7 @@
-
+