diff --git a/.gitattributes b/.gitattributes index 44e16cf1d..7d89f37b2 100644 --- a/.gitattributes +++ b/.gitattributes @@ -11,10 +11,11 @@ # 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. - *.memmap filter=lfs diff=lfs merge=lfs -text *.stl filter=lfs diff=lfs merge=lfs -text *.safetensors filter=lfs diff=lfs merge=lfs -text *.mp4 filter=lfs diff=lfs merge=lfs -text *.arrow filter=lfs diff=lfs merge=lfs -text *.json !text !filter !merge !diff +tests/artifacts/cameras/*.png filter=lfs diff=lfs merge=lfs -text +*.bag filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore index d6c51c90d..4ab886933 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +# Dev scripts +.dev + # Logging logs tmp @@ -26,6 +29,7 @@ outputs # VS Code .vscode +.devcontainer # HPC nautilus/*.yaml @@ -91,10 +95,8 @@ coverage.xml .hypothesis/ .pytest_cache/ -# Ignore .cache except calibration +# Ignore .cache .cache/* -!.cache/calibration/ -!.cache/calibration/** # Translations *.mo diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a778ce0e9..e1f971d39 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -37,18 +37,18 @@ repos: - id: trailing-whitespace - repo: https://github.com/adhtruong/mirrors-typos - rev: v1.31.1 + rev: v1.33.1 hooks: - id: typos args: [--force-exclude] - repo: https://github.com/asottile/pyupgrade - rev: v3.19.1 + rev: v3.20.0 hooks: - id: pyupgrade - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.11.5 + rev: v0.11.13 hooks: - id: ruff args: [--fix] @@ -57,12 +57,12 @@ repos: ##### Security ##### - repo: https://github.com/gitleaks/gitleaks - rev: v8.24.3 + rev: v8.27.2 hooks: - id: gitleaks - repo: https://github.com/woodruffw/zizmor-pre-commit - rev: v1.5.2 + rev: v1.9.0 hooks: - id: zizmor diff --git a/Makefile b/Makefile index c82483cc3..9457dbe6e 100644 --- a/Makefile +++ b/Makefile @@ -40,6 +40,8 @@ test-end-to-end: ${MAKE} DEVICE=$(DEVICE) test-diffusion-ete-eval ${MAKE} DEVICE=$(DEVICE) test-tdmpc-ete-train ${MAKE} DEVICE=$(DEVICE) test-tdmpc-ete-eval + ${MAKE} DEVICE=$(DEVICE) test-smolvla-ete-train + ${MAKE} DEVICE=$(DEVICE) test-smolvla-ete-eval test-act-ete-train: python lerobot/scripts/train.py \ @@ -48,6 +50,7 @@ test-act-ete-train: --policy.n_action_steps=20 \ --policy.chunk_size=20 \ --policy.device=$(DEVICE) \ + --policy.push_to_hub=false \ --env.type=aloha \ --env.episode_length=5 \ --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \ @@ -85,6 +88,7 @@ test-diffusion-ete-train: --policy.diffusion_step_embed_dim=32 \ --policy.num_inference_steps=10 \ --policy.device=$(DEVICE) \ + --policy.push_to_hub=false \ --env.type=pusht \ --env.episode_length=5 \ --dataset.repo_id=lerobot/pusht \ @@ -114,6 +118,7 @@ test-tdmpc-ete-train: python lerobot/scripts/train.py \ --policy.type=tdmpc \ --policy.device=$(DEVICE) \ + --policy.push_to_hub=false \ --env.type=xarm \ --env.task=XarmLift-v0 \ --env.episode_length=5 \ @@ -140,3 +145,36 @@ test-tdmpc-ete-eval: --env.task=XarmLift-v0 \ --eval.n_episodes=1 \ --eval.batch_size=1 + + +test-smolvla-ete-train: + python lerobot/scripts/train.py \ + --policy.type=smolvla \ + --policy.n_action_steps=20 \ + --policy.chunk_size=20 \ + --policy.device=$(DEVICE) \ + --policy.push_to_hub=false \ + --env.type=aloha \ + --env.episode_length=5 \ + --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \ + --dataset.image_transforms.enable=true \ + --dataset.episodes="[0]" \ + --batch_size=2 \ + --steps=4 \ + --eval_freq=2 \ + --eval.n_episodes=1 \ + --eval.batch_size=1 \ + --save_freq=2 \ + --save_checkpoint=true \ + --log_freq=1 \ + --wandb.enable=false \ + --output_dir=tests/outputs/smolvla/ + +test-smolvla-ete-eval: + python lerobot/scripts/eval.py \ + --policy.path=tests/outputs/smolvla/checkpoints/000004/pretrained_model \ + --policy.device=$(DEVICE) \ + --env.type=aloha \ + --env.episode_length=5 \ + --eval.n_episodes=1 \ + --eval.batch_size=1 diff --git a/README.md b/README.md index dd55dbff3..153a3a215 100644 --- a/README.md +++ b/README.md @@ -23,22 +23,36 @@

-

- Build Your Own SO-100 Robot!

+

+ Build Your Own SO-101 Robot!

- SO-100 leader and follower arms +
+ SO-101 follower arm + SO-101 leader arm +
-

Meet the SO-100 – Just $110 per arm!

+ +

Meet the updated SO100, the SO-101 – Just €114 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.

+

+ See the full SO-101 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.

+

Want to take it to the next level? Make your SO-101 mobile by building LeKiwi!

+

Check out the LeKiwi tutorial and bring your robot to life on wheels.

LeKiwi mobile robot
@@ -51,7 +65,6 @@ --- - 🤗 LeRobot aims to provide models, datasets, and tools for real-world robotics in PyTorch. The goal is to lower the barrier to entry to robotics so that everyone can contribute and benefit from sharing datasets and pretrained models. 🤗 LeRobot contains state-of-the-art approaches that have been shown to transfer to the real-world with a focus on imitation learning and reinforcement learning. @@ -77,6 +90,7 @@ ### Acknowledgment +- The LeRobot team 🤗 for building SmolVLA [Paper](https://arxiv.org/abs/2506.01844), [Blog](https://huggingface.co/blog/smolvla). - Thanks to Tony Zhao, Zipeng Fu and colleagues for open sourcing ACT policy, ALOHA environments and datasets. Ours are adapted from [ALOHA](https://tonyzhaozh.github.io/aloha) and [Mobile ALOHA](https://mobile-aloha.github.io). - Thanks to Cheng Chi, Zhenjia Xu and colleagues for open sourcing Diffusion policy, Pusht environment and datasets, as well as UMI datasets. Ours are adapted from [Diffusion Policy](https://diffusion-policy.cs.columbia.edu) and [UMI Gripper](https://umi-gripper.github.io). - Thanks to Nicklas Hansen, Yunhai Feng and colleagues for open sourcing TDMPC policy, Simxarm environments and datasets. Ours are adapted from [TDMPC](https://github.com/nicklashansen/tdmpc) and [FOWM](https://www.yunhaifeng.com/FOWM). @@ -116,7 +130,7 @@ pip install -e . ``` > **NOTE:** If you encounter build errors, you may need to install additional dependencies (`cmake`, `build-essential`, and `ffmpeg libs`). On Linux, run: -`sudo apt-get install cmake build-essential python3-dev pkg-config libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev libswscale-dev libswresample-dev libavfilter-dev pkg-config`. For other systems, see: [Compiling PyAV](https://pyav.org/docs/develop/overview/installation.html#bring-your-own-ffmpeg) +`sudo apt-get install cmake build-essential python3-dev pkg-config libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev libswscale-dev libswresample-dev libavfilter-dev`. For other systems, see: [Compiling PyAV](https://pyav.org/docs/develop/overview/installation.html#bring-your-own-ffmpeg) For simulations, 🤗 LeRobot comes with gymnasium environments that can be installed as extras: - [aloha](https://github.com/huggingface/gym-aloha) @@ -174,7 +188,6 @@ Under the hood, the `LeRobotDataset` format makes use of several ways to seriali Here are the important details and internal structure organization of a typical `LeRobotDataset` instantiated with `dataset = LeRobotDataset("lerobot/aloha_static_coffee")`. The exact features will change from dataset to dataset but not the main aspects: ``` -TODO: IMPROVE dataset attributes: ├ hf_dataset: a Hugging Face dataset (backed by Arrow/parquet). Typical features example: │ ├ observation.images.cam_high (VideoFrame): @@ -185,9 +198,9 @@ dataset attributes: │ ├ episode_index (int64): index of the episode for this sample │ ├ frame_index (int64): index of the frame for this sample in the episode ; starts at 0 for each episode │ ├ timestamp (float32): timestamp in the episode - │ ├ next.done (bool): indicates the end of en episode ; True for the last frame in each episode + │ ├ next.done (bool): indicates the end of an episode ; True for the last frame in each episode │ └ index (int64): general index in the whole dataset - ├ meta: contains 2 tensors with the start and end indices of each episode + ├ episode_data_index: contains 2 tensors with the start and end indices of each episode │ ├ from (1D int64 tensor): first frame index for each episode — shape (num episodes,) starts with 0 │ └ to: (1D int64 tensor): last frame index for each episode — shape (num episodes,) ├ stats: a dictionary of statistics (max, mean, min, std) for each feature in the dataset, for instance @@ -234,7 +247,7 @@ See `python -m lerobot.scripts.eval --help` for more instructions. ### Train your own policy -Check out [example 3](./examples/3_train_policy.py) that illustrate how to train a model using our core library in python, and [example 4](./examples/4_train_policy_with_script.md) that shows how to use our training script from command line. +Check out [example 3](./examples/3_train_policy.py) that illustrates how to train a model using our core library in python, and [example 4](./examples/4_train_policy_with_script.md) that shows how to use our training script from command line. To use wandb for logging training and evaluation curves, make sure you've run `wandb login` as a one-time setup step. Then, when running the training command above, enable WandB in the configuration by adding `--wandb.enable=true`. @@ -285,7 +298,7 @@ Once you have trained a policy you may upload it to the Hugging Face hub using a You first need to find the checkpoint folder located inside your experiment directory (e.g. `outputs/train/2024-05-05/20-21-12_aloha_act_default/checkpoints/002500`). Within that there is a `pretrained_model` directory which should contain: - `config.json`: A serialized version of the policy configuration (following the policy's dataclass config). - `model.safetensors`: A set of `torch.nn.Module` parameters, saved in [Hugging Face Safetensors](https://huggingface.co/docs/safetensors/index) format. -- `train_config.json`: A consolidated configuration containing all parameter userd for training. The policy configuration should match `config.json` exactly. Thisis useful for anyone who wants to evaluate your policy or for reproducibility. +- `train_config.json`: A consolidated configuration containing all parameters used for training. The policy configuration should match `config.json` exactly. This is useful for anyone who wants to evaluate your policy or for reproducibility. To upload these to the hub, run the following: ```bash @@ -324,7 +337,7 @@ with profile( If you want, you can cite this work with: ```bibtex @misc{cadene2024lerobot, - author = {Cadene, Remi and Alibert, Simon and Soare, Alexander and Gallouedec, Quentin and Zouitine, Adil and Wolf, Thomas}, + author = {Cadene, Remi and Alibert, Simon and Soare, Alexander and Gallouedec, Quentin and Zouitine, Adil and Palma, Steven and Kooijmans, Pepijn and Aractingi, Michel and Shukor, Mustafa and Aubakirova, Dana and Russi, Martino and Capuano, Francesco and Pascale, Caroline and Choghari, Jade and Moss, Jess and Wolf, Thomas}, title = {LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch}, howpublished = "\url{https://github.com/huggingface/lerobot}", year = {2024} @@ -332,6 +345,15 @@ If you want, you can cite this work with: ``` Additionally, if you are using any of the particular policy architecture, pretrained models, or datasets, it is recommended to cite the original authors of the work as they appear below: +- [SmolVLA](https://arxiv.org/abs/2506.01844) +```bibtex +@article{shukor2025smolvla, + title={SmolVLA: A Vision-Language-Action Model for Affordable and Efficient Robotics}, + author={Shukor, Mustafa and Aubakirova, Dana and Capuano, Francesco and Kooijmans, Pepijn and Palma, Steven and Zouitine, Adil and Aractingi, Michel and Pascal, Caroline and Russi, Martino and Marafioti, Andres and Alibert, Simon and Cord, Matthieu and Wolf, Thomas and Cadene, Remi}, + journal={arXiv preprint arXiv:2506.01844}, + year={2025} +} +``` - [Diffusion Policy](https://diffusion-policy.cs.columbia.edu) ```bibtex @@ -372,6 +394,19 @@ Additionally, if you are using any of the particular policy architecture, pretra year={2024} } ``` + + +- [HIL-SERL](https://hil-serl.github.io/) +```bibtex +@Article{luo2024hilserl, +title={Precise and Dexterous Robotic Manipulation via Human-in-the-Loop Reinforcement Learning}, +author={Jianlan Luo and Charles Xu and Jeffrey Wu and Sergey Levine}, +year={2024}, +eprint={2410.21845}, +archivePrefix={arXiv}, +primaryClass={cs.RO} +} +``` ## Star History [![Star History Chart](https://api.star-history.com/svg?repos=huggingface/lerobot&type=Timeline)](https://star-history.com/#huggingface/lerobot&Timeline) diff --git a/benchmarks/video/capture_camera_feed.py b/benchmarks/video/capture_camera_feed.py index ce248f20b..8f8530532 100755 --- a/benchmarks/video/capture_camera_feed.py +++ b/benchmarks/video/capture_camera_feed.py @@ -55,7 +55,7 @@ def display_and_save_video_stream(output_dir: Path, fps: int, width: int, height if not ret: print("Error: Could not read frame.") break - rr.log("video/stream", rr.Image(frame.numpy()), static=True) + rr.log("video/stream", rr.Image(frame), static=True) cv2.imwrite(str(capture_dir / f"frame_{frame_index:06d}.png"), frame) frame_index += 1 diff --git a/examples/3_train_policy.py b/examples/3_train_policy.py index ded61a071..f2de79db8 100644 --- a/examples/3_train_policy.py +++ b/examples/3_train_policy.py @@ -12,7 +12,7 @@ # 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. +"""This script 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 examples/2_evaluate_pretrained_policy.py diff --git a/examples/4_train_policy_with_script.md b/examples/4_train_policy_with_script.md index b23d22713..ff8913016 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 `--policy.device=cpu` (`--policy.device=mps` respectively). However, be advised that the code executes much slower on cpu. +> **Note:** The following assumes 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 @@ -23,7 +23,7 @@ def train(cfg: TrainPipelineConfig): You can inspect the `TrainPipelineConfig` defined in [`lerobot/configs/train.py`](../../lerobot/configs/train.py) (which is heavily commented and meant to be a reference to understand any option) -When running the script, inputs for the command line are parsed thanks to the `@parser.wrap()` decorator and an instance of this class is automatically generated. Under the hood, this is done with [Draccus](https://github.com/dlwh/draccus) which is a tool dedicated for this purpose. If you're familiar with Hydra, Draccus can similarly load configurations from config files (.json, .yaml) and also override their values through command line inputs. Unlike Hydra, these configurations are pre-defined in the code through dataclasses rather than being defined entirely in config files. This allows for more rigorous serialization/deserialization, typing, and to manipulate configuration as objects directly in the code and not as dictionaries or namespaces (which enables nice features in an IDE such as autocomplete, jump-to-def, etc.) +When running the script, inputs for the command line are parsed thanks to the `@parser.wrap()` decorator and an instance of this class is automatically generated. Under the hood, this is done with [Draccus](https://github.com/dlwh/draccus) which is a tool dedicated to this purpose. If you're familiar with Hydra, Draccus can similarly load configurations from config files (.json, .yaml) and also override their values through command line inputs. Unlike Hydra, these configurations are pre-defined in the code through dataclasses rather than being defined entirely in config files. This allows for more rigorous serialization/deserialization, typing, and to manipulate configuration as objects directly in the code and not as dictionaries or namespaces (which enables nice features in an IDE such as autocomplete, jump-to-def, etc.) Let's have a look at a simplified example. Amongst other attributes, the training config has the following attributes: ```python @@ -43,7 +43,7 @@ class DatasetConfig: ``` This creates a hierarchical relationship where, for example assuming we have a `cfg` instance of `TrainPipelineConfig`, we can access the `repo_id` value with `cfg.dataset.repo_id`. -From the command line, we can specify this value with using a very similar syntax `--dataset.repo_id=repo/id`. +From the command line, we can specify this value by using a very similar syntax `--dataset.repo_id=repo/id`. By default, every field takes its default value specified in the dataclass. If a field doesn't have a default value, it needs to be specified either from the command line or from a config file – which path is also given in the command line (more in this below). In the example above, the `dataset` field doesn't have a default value which means it must be specified. @@ -135,7 +135,7 @@ will start a training run with the same configuration used for training [lerobot ## Resume training -Being able to resume a training run is important in case it crashed or aborted for any reason. We'll demonstrate how to that here. +Being able to resume a training run is important in case it crashed or aborted for any reason. We'll demonstrate how to do that here. Let's reuse the command from the previous run and add a few more options: ```bash diff --git a/lerobot/common/policies/pi0fast/modeling_pi0fast.py b/lerobot/common/policies/pi0fast/modeling_pi0fast.py index 36aafce94..dbf5266b1 100644 --- a/lerobot/common/policies/pi0fast/modeling_pi0fast.py +++ b/lerobot/common/policies/pi0fast/modeling_pi0fast.py @@ -17,7 +17,7 @@ """ π0+FAST: Efficient Action Tokenization for Vision-Language-Action Models -[Paper](https://arxiv.org/abs/2501.09747) +[Paper](https://huggingface.co/papers/2501.09747) [Jax code](https://github.com/Physical-Intelligence/openpi) Designed by Physical Intelligence. Ported from Jax by Hugging Face. @@ -56,7 +56,7 @@ from transformers import AutoProcessor, AutoTokenizer, PaliGemmaForConditionalGe from transformers.cache_utils import HybridCache, StaticCache from transformers.models.auto import CONFIG_MAPPING -from lerobot.common.constants import ACTION, OBS_ROBOT +from lerobot.common.constants import ACTION, OBS_STATE from lerobot.common.policies.normalize import Normalize, Unnormalize from lerobot.common.policies.pi0fast.configuration_pi0fast import PI0FASTConfig from lerobot.common.policies.pretrained import PreTrainedPolicy @@ -192,6 +192,11 @@ class PI0FASTPolicy(PreTrainedPolicy): actions[:, :, motor_idx] = aloha_gripper_from_angular_inv(actions[:, :, motor_idx]) return actions + @torch.no_grad + def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor: + """Predict a chunk of actions given environment observations.""" + raise NotImplementedError("Currently not implemented for PI0FAST") + @torch.no_grad def select_action(self, batch: dict[str, Tensor]) -> Tensor: """Select a single action given environment observations. @@ -203,7 +208,7 @@ class PI0FASTPolicy(PreTrainedPolicy): self.eval() if self.config.adapt_to_pi_aloha: - batch[OBS_ROBOT] = self._pi_aloha_decode_state(batch[OBS_ROBOT]) + batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE]) batch = self.normalize_inputs(batch) @@ -231,7 +236,7 @@ class PI0FASTPolicy(PreTrainedPolicy): def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]: if self.config.adapt_to_pi_aloha: - batch[OBS_ROBOT] = self._pi_aloha_decode_state(batch[OBS_ROBOT]) + batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE]) batch[ACTION] = self._pi_aloha_encode_actions_inv(batch[ACTION]) batch = self.normalize_inputs(batch) batch = self.normalize_targets(batch) @@ -516,7 +521,7 @@ class PI0FAST(nn.Module): interpolate_like_pi=self.config.interpolate_like_pi, ) - # Normalize from range [0,1] to [-1,1] as expacted by siglip + # Normalize from range [0,1] to [-1,1] as expected by siglip img = img * 2.0 - 1.0 bsize = img.shape[0] @@ -677,12 +682,12 @@ class PI0FAST(nn.Module): return new_tokens, new_ar_masks, new_padding_mask, new_loss_mask, new_targets, new_token_type_ids def forward(self, batch: dict[str, Tensor]): - device = batch[OBS_ROBOT].device + device = batch[OBS_STATE].device # TODO: keep like this or move to the policy .forward images, img_masks = self.prepare_images(batch) padded_outs = self.create_input_tokens( - state=batch[OBS_ROBOT], + state=batch[OBS_STATE], lang_text=batch["task"], actions=batch[ACTION], ) @@ -849,7 +854,7 @@ class PI0FAST(nn.Module): # TODO: keep like this or move to the policy .forward images, img_masks = self.prepare_images(batch) - padded_outs = self.create_input_tokens(state=batch[OBS_ROBOT], lang_text=batch["task"], actions=None) + padded_outs = self.create_input_tokens(state=batch[OBS_STATE], lang_text=batch["task"], actions=None) embs, pad_masks, att_masks2, targets, loss_mask, token_type_ids = self.embed_inputs( images, img_masks, @@ -878,7 +883,11 @@ class PI0FAST(nn.Module): return actions def embed_image(self, image: torch.Tensor): - return self.pi0_paligemma.get_image_features(image) + # Handle different transformers versions + if hasattr(self.pi0_paligemma, "get_image_features"): + return self.pi0_paligemma.get_image_features(image) + else: + return self.pi0_paligemma.model.get_image_features(image) def embed_inputs( self, diff --git a/pyproject.toml b/pyproject.toml index bbed7a9b4..3686532ed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,11 +49,11 @@ dependencies = [ "datasets>=2.19.0", "deepdiff>=7.0.1", "diffusers>=0.27.2", - "draccus>=0.10.0", + "draccus==0.10.0", "einops>=0.8.0", "flask>=3.0.3", "gdown>=5.1.0", - "gymnasium==0.29.1", # TODO(rcadene, aliberts): Make gym 1.0.0 work + "gymnasium==0.29.1", # TODO(rcadene, aliberts): Make gym 1.0.0 work "h5py>=3.10.0", "huggingface-hub[hf-transfer,cli]>=0.27.1 ; python_version < '4.0'", "imageio[ffmpeg]>=2.34.0", @@ -62,11 +62,13 @@ dependencies = [ "omegaconf>=2.3.0", "opencv-python-headless>=4.9.0", "packaging>=24.2", - "av>=12.0.5", - "pymunk>=6.6.0", + "av>=14.2.0", + "pymunk>=6.6.0,<7.0.0", "pynput>=1.7.7", + "pyserial>=3.5", "pyzmq>=26.2.1", "rerun-sdk>=0.21.0", + "scipy>=1.14.0", "termcolor>=2.4.0", "torch>=2.2.1", "torchcodec>=0.2.1; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')", @@ -77,22 +79,28 @@ dependencies = [ [project.optional-dependencies] aloha = ["gym-aloha>=0.1.1 ; python_version < '4.0'"] +docs = ["hf-doc-builder @ git+https://github.com/huggingface/doc-builder.git@main", "watchdog >= 6.0.0"] dev = ["pre-commit>=3.7.0", "debugpy>=1.8.1"] dora = [ "gym-dora @ git+https://github.com/dora-rs/dora-lerobot.git#subdirectory=gym_dora ; python_version < '4.0'", ] -dynamixel = ["dynamixel-sdk>=3.7.31", "pynput>=1.7.7"] -feetech = ["feetech-servo-sdk>=1.0.0", "pynput>=1.7.7"] -intelrealsense = ["pyrealsense2>=2.55.1.6486 ; sys_platform != 'darwin'"] -pi0 = ["transformers>=4.48.0"] +dynamixel = ["dynamixel-sdk>=3.7.31"] +feetech = ["feetech-servo-sdk>=1.0.0"] +gamepad = ["pygame>=2.5.1", "hidapi>=0.14.0"] +intelrealsense = [ + "pyrealsense2>=2.55.1.6486 ; sys_platform != 'darwin'", + "pyrealsense2-macosx>=2.54 ; sys_platform == 'darwin'", +] +pi0 = ["transformers>=4.50.3"] +smolvla = ["transformers>=4.50.3", "num2words>=0.5.14", "accelerate>=1.7.0", "safetensors>=0.4.3"] pusht = ["gym-pusht>=0.1.5 ; python_version < '4.0'"] stretch = [ "hello-robot-stretch-body>=0.7.27 ; python_version < '4.0' and sys_platform == 'linux'", "pyrender @ git+https://github.com/mmatl/pyrender.git ; sys_platform == 'linux'", - "pyrealsense2>=2.55.1.6486 ; sys_platform != 'darwin'", - "pynput>=1.7.7", + "pyrealsense2>=2.55.1.6486 ; sys_platform != 'darwin'" ] -test = ["pytest>=8.1.0", "pytest-cov>=5.0.0", "pyserial>=3.5"] +test = ["pytest>=8.1.0", "pytest-timeout>=2.4.0", "pytest-cov>=5.0.0", "pyserial>=3.5", "mock-serial>=0.0.1 ; sys_platform != 'win32'"] +hilserl = ["transformers>=4.50.3", "gym-hil>=0.1.8", "protobuf>=5.29.3", "grpcio==1.71.0"] umi = ["imagecodecs>=2024.1.1"] video_benchmark = ["scikit-image>=0.23.2", "pandas>=2.2.2"] xarm = ["gym-xarm>=0.1.1 ; python_version < '4.0'"] @@ -103,11 +111,14 @@ requires-poetry = ">=2.1" [tool.ruff] line-length = 110 target-version = "py310" -exclude = ["tests/artifacts/**/*.safetensors"] +exclude = ["tests/artifacts/**/*.safetensors", "*_pb2.py", "*_pb2_grpc.py"] [tool.ruff.lint] select = ["E4", "E7", "E9", "F", "I", "N", "B", "C4", "SIM"] +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["F401", "F403"] + [tool.bandit] exclude_dirs = [ "tests", diff --git a/tests/artifacts/policies/save_policy_to_safetensors.py b/tests/artifacts/policies/save_policy_to_safetensors.py index 6573782e7..6ccb47c3e 100644 --- a/tests/artifacts/policies/save_policy_to_safetensors.py +++ b/tests/artifacts/policies/save_policy_to_safetensors.py @@ -32,7 +32,7 @@ def get_policy_stats(ds_repo_id: str, policy_name: str, policy_kwargs: dict): train_cfg = TrainPipelineConfig( # TODO(rcadene, aliberts): remove dataset download dataset=DatasetConfig(repo_id=ds_repo_id, episodes=[0]), - policy=make_policy_config(policy_name, **policy_kwargs), + policy=make_policy_config(policy_name, push_to_hub=False, **policy_kwargs), ) train_cfg.validate() # Needed for auto-setting some parameters diff --git a/tests/conftest.py b/tests/conftest.py index 7eec94bf8..69dd3049b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -19,9 +19,7 @@ import traceback import pytest from serial import SerialException -from lerobot import available_cameras, available_motors, available_robots -from lerobot.common.robot_devices.robots.utils import make_robot -from tests.utils import DEVICE, make_camera, make_motors_bus +from tests.utils import DEVICE # Import fixture modules as plugins pytest_plugins = [ @@ -64,21 +62,6 @@ def _check_component_availability(component_type, available_components, make_com return False -@pytest.fixture -def is_robot_available(robot_type): - return _check_component_availability(robot_type, available_robots, make_robot) - - -@pytest.fixture -def is_camera_available(camera_type): - return _check_component_availability(camera_type, available_cameras, make_camera) - - -@pytest.fixture -def is_motor_available(motor_type): - return _check_component_availability(motor_type, available_motors, make_motors_bus) - - @pytest.fixture def patch_builtins_input(monkeypatch): def print_text(text=None): diff --git a/tests/optim/test_optimizers.py b/tests/optim/test_optimizers.py index 997e14fe9..630353fca 100644 --- a/tests/optim/test_optimizers.py +++ b/tests/optim/test_optimizers.py @@ -21,6 +21,7 @@ from lerobot.common.constants import ( from lerobot.common.optim.optimizers import ( AdamConfig, AdamWConfig, + MultiAdamConfig, SGDConfig, load_optimizer_state, save_optimizer_state, @@ -33,13 +34,21 @@ from lerobot.common.optim.optimizers import ( (AdamConfig, torch.optim.Adam), (AdamWConfig, torch.optim.AdamW), (SGDConfig, torch.optim.SGD), + (MultiAdamConfig, dict), ], ) def test_optimizer_build(config_cls, expected_class, model_params): config = config_cls() - optimizer = config.build(model_params) - assert isinstance(optimizer, expected_class) - assert optimizer.defaults["lr"] == config.lr + if config_cls == MultiAdamConfig: + params_dict = {"default": model_params} + optimizer = config.build(params_dict) + assert isinstance(optimizer, expected_class) + assert isinstance(optimizer["default"], torch.optim.Adam) + assert optimizer["default"].defaults["lr"] == config.lr + else: + optimizer = config.build(model_params) + assert isinstance(optimizer, expected_class) + assert optimizer.defaults["lr"] == config.lr def test_save_optimizer_state(optimizer, tmp_path): @@ -54,3 +63,180 @@ def test_save_and_load_optimizer_state(model_params, optimizer, tmp_path): loaded_optimizer = load_optimizer_state(loaded_optimizer, tmp_path) torch.testing.assert_close(optimizer.state_dict(), loaded_optimizer.state_dict()) + + +@pytest.fixture +def base_params_dict(): + return { + "actor": [torch.nn.Parameter(torch.randn(10, 10))], + "critic": [torch.nn.Parameter(torch.randn(5, 5))], + "temperature": [torch.nn.Parameter(torch.randn(3, 3))], + } + + +@pytest.mark.parametrize( + "config_params, expected_values", + [ + # Test 1: Basic configuration with different learning rates + ( + { + "lr": 1e-3, + "weight_decay": 1e-4, + "optimizer_groups": { + "actor": {"lr": 1e-4}, + "critic": {"lr": 5e-4}, + "temperature": {"lr": 2e-3}, + }, + }, + { + "actor": {"lr": 1e-4, "weight_decay": 1e-4, "betas": (0.9, 0.999)}, + "critic": {"lr": 5e-4, "weight_decay": 1e-4, "betas": (0.9, 0.999)}, + "temperature": {"lr": 2e-3, "weight_decay": 1e-4, "betas": (0.9, 0.999)}, + }, + ), + # Test 2: Different weight decays and beta values + ( + { + "lr": 1e-3, + "weight_decay": 1e-4, + "optimizer_groups": { + "actor": {"lr": 1e-4, "weight_decay": 1e-5}, + "critic": {"lr": 5e-4, "weight_decay": 1e-6}, + "temperature": {"lr": 2e-3, "betas": (0.95, 0.999)}, + }, + }, + { + "actor": {"lr": 1e-4, "weight_decay": 1e-5, "betas": (0.9, 0.999)}, + "critic": {"lr": 5e-4, "weight_decay": 1e-6, "betas": (0.9, 0.999)}, + "temperature": {"lr": 2e-3, "weight_decay": 1e-4, "betas": (0.95, 0.999)}, + }, + ), + # Test 3: Epsilon parameter customization + ( + { + "lr": 1e-3, + "weight_decay": 1e-4, + "optimizer_groups": { + "actor": {"lr": 1e-4, "eps": 1e-6}, + "critic": {"lr": 5e-4, "eps": 1e-7}, + "temperature": {"lr": 2e-3, "eps": 1e-8}, + }, + }, + { + "actor": {"lr": 1e-4, "weight_decay": 1e-4, "betas": (0.9, 0.999), "eps": 1e-6}, + "critic": {"lr": 5e-4, "weight_decay": 1e-4, "betas": (0.9, 0.999), "eps": 1e-7}, + "temperature": {"lr": 2e-3, "weight_decay": 1e-4, "betas": (0.9, 0.999), "eps": 1e-8}, + }, + ), + ], +) +def test_multi_adam_configuration(base_params_dict, config_params, expected_values): + # Create config with the given parameters + config = MultiAdamConfig(**config_params) + optimizers = config.build(base_params_dict) + + # Verify optimizer count and keys + assert len(optimizers) == len(expected_values) + assert set(optimizers.keys()) == set(expected_values.keys()) + + # Check that all optimizers are Adam instances + for opt in optimizers.values(): + assert isinstance(opt, torch.optim.Adam) + + # Verify hyperparameters for each optimizer + for name, expected in expected_values.items(): + optimizer = optimizers[name] + for param, value in expected.items(): + assert optimizer.defaults[param] == value + + +@pytest.fixture +def multi_optimizers(base_params_dict): + config = MultiAdamConfig( + lr=1e-3, + optimizer_groups={ + "actor": {"lr": 1e-4}, + "critic": {"lr": 5e-4}, + "temperature": {"lr": 2e-3}, + }, + ) + return config.build(base_params_dict) + + +def test_save_multi_optimizer_state(multi_optimizers, tmp_path): + # Save optimizer states + save_optimizer_state(multi_optimizers, tmp_path) + + # Verify that directories were created for each optimizer + for name in multi_optimizers: + assert (tmp_path / name).is_dir() + assert (tmp_path / name / OPTIMIZER_STATE).is_file() + assert (tmp_path / name / OPTIMIZER_PARAM_GROUPS).is_file() + + +def test_save_and_load_multi_optimizer_state(base_params_dict, multi_optimizers, tmp_path): + # Option 1: Add a minimal backward pass to populate optimizer states + for name, params in base_params_dict.items(): + if name in multi_optimizers: + # Create a dummy loss and do backward + dummy_loss = params[0].sum() + dummy_loss.backward() + # Perform an optimization step + multi_optimizers[name].step() + # Zero gradients for next steps + multi_optimizers[name].zero_grad() + + # Save optimizer states + save_optimizer_state(multi_optimizers, tmp_path) + + # Create new optimizers with the same config + config = MultiAdamConfig( + lr=1e-3, + optimizer_groups={ + "actor": {"lr": 1e-4}, + "critic": {"lr": 5e-4}, + "temperature": {"lr": 2e-3}, + }, + ) + new_optimizers = config.build(base_params_dict) + + # Load optimizer states + loaded_optimizers = load_optimizer_state(new_optimizers, tmp_path) + + # Verify state dictionaries match + for name in multi_optimizers: + torch.testing.assert_close(multi_optimizers[name].state_dict(), loaded_optimizers[name].state_dict()) + + +def test_save_and_load_empty_multi_optimizer_state(base_params_dict, tmp_path): + """Test saving and loading optimizer states even when the state is empty (no backward pass).""" + # Create config and build optimizers + config = MultiAdamConfig( + lr=1e-3, + optimizer_groups={ + "actor": {"lr": 1e-4}, + "critic": {"lr": 5e-4}, + "temperature": {"lr": 2e-3}, + }, + ) + optimizers = config.build(base_params_dict) + + # Save optimizer states without any backward pass (empty state) + save_optimizer_state(optimizers, tmp_path) + + # Create new optimizers with the same config + new_optimizers = config.build(base_params_dict) + + # Load optimizer states + loaded_optimizers = load_optimizer_state(new_optimizers, tmp_path) + + # Verify hyperparameters match even with empty state + for name, optimizer in optimizers.items(): + assert optimizer.defaults["lr"] == loaded_optimizers[name].defaults["lr"] + assert optimizer.defaults["weight_decay"] == loaded_optimizers[name].defaults["weight_decay"] + assert optimizer.defaults["betas"] == loaded_optimizers[name].defaults["betas"] + + # Verify state dictionaries match (they will be empty) + torch.testing.assert_close( + optimizer.state_dict()["param_groups"], loaded_optimizers[name].state_dict()["param_groups"] + )