Merge 38e36d2d4c into b43ece8934

Co-authored-by: makolon <smakolon385@gmail.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>

README.md

@@ -116,7 +116,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 python-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 pkg-config`. 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)

docker/lerobot-gpu-dev/Dockerfile

@@ -14,7 +14,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     tcpdump sysstat screen tmux \
     libglib2.0-0 libgl1-mesa-glx libegl1-mesa \
     speech-dispatcher portaudio19-dev libgeos-dev \
-    python${PYTHON_VERSION} python${PYTHON_VERSION}-venv \
+    python${PYTHON_VERSION} python${PYTHON_VERSION}-venv python${PYTHON_VERSION}-dev \
     && apt-get clean && rm -rf /var/lib/apt/lists/*
 
 # Install ffmpeg build dependencies. See:

examples/4_train_policy_with_script.md

@@ -4,7 +4,7 @@ This tutorial will explain the training script, how to use it, and particularly
 
 ## The training script
 
-LeRobot offers a training script at [`lerobot/scripts/train.py`](../../lerobot/scripts/train.py). At a high level it does the following:
+LeRobot offers a training script at [`lerobot/scripts/train.py`](../lerobot/scripts/train.py). At a high level it does the following:
 
 - Initialize/load a configuration for the following steps using.
 - Instantiates a dataset.
@@ -21,7 +21,7 @@ In the training script, the main function `train` expects a `TrainPipelineConfig
 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)
+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.)
 
@@ -50,7 +50,7 @@ By default, every field takes its default value specified in the dataclass. If a
 
 ## Specifying values from the CLI
 
-Let's say that we want to train [Diffusion Policy](../../lerobot/common/policies/diffusion) on the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using the [gym_pusht](https://github.com/huggingface/gym-pusht) environment for evaluation. The command to do so would look like this:
+Let's say that we want to train [Diffusion Policy](../lerobot/common/policies/diffusion) on the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using the [gym_pusht](https://github.com/huggingface/gym-pusht) environment for evaluation. The command to do so would look like this:
 ```bash
 python lerobot/scripts/train.py \
     --dataset.repo_id=lerobot/pusht \
@@ -60,10 +60,10 @@ python lerobot/scripts/train.py \
 
 Let's break this down:
 - To specify the dataset, we just need to specify its `repo_id` on the hub which is the only required argument in the `DatasetConfig`. The rest of the fields have default values and in this case we are fine with those so we can just add the option `--dataset.repo_id=lerobot/pusht`.
-- To specify the policy, we can just select diffusion policy using `--policy` appended with `.type`. Here, `.type` is a special argument which allows us to select config classes inheriting from `draccus.ChoiceRegistry` and that have been decorated with the `register_subclass()` method. To have a better explanation of this feature, have a look at this [Draccus demo](https://github.com/dlwh/draccus?tab=readme-ov-file#more-flexible-configuration-with-choice-types). In our code, we use this mechanism mainly to select policies, environments, robots, and some other components like optimizers. The policies available to select are located in [lerobot/common/policies](../../lerobot/common/policies)
-- Similarly, we select the environment with `--env.type=pusht`. The different environment configs are available in [`lerobot/common/envs/configs.py`](../../lerobot/common/envs/configs.py)
+- To specify the policy, we can just select diffusion policy using `--policy` appended with `.type`. Here, `.type` is a special argument which allows us to select config classes inheriting from `draccus.ChoiceRegistry` and that have been decorated with the `register_subclass()` method. To have a better explanation of this feature, have a look at this [Draccus demo](https://github.com/dlwh/draccus?tab=readme-ov-file#more-flexible-configuration-with-choice-types). In our code, we use this mechanism mainly to select policies, environments, robots, and some other components like optimizers. The policies available to select are located in [lerobot/common/policies](../lerobot/common/policies)
+- Similarly, we select the environment with `--env.type=pusht`. The different environment configs are available in [`lerobot/common/envs/configs.py`](../lerobot/common/envs/configs.py)
 
-Let's see another example. Let's say you've been training [ACT](../../lerobot/common/policies/act) on [lerobot/aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human) using the [gym-aloha](https://github.com/huggingface/gym-aloha) environment for evaluation with:
+Let's see another example. Let's say you've been training [ACT](../lerobot/common/policies/act) on [lerobot/aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human) using the [gym-aloha](https://github.com/huggingface/gym-aloha) environment for evaluation with:
 ```bash
 python lerobot/scripts/train.py \
     --policy.type=act \
@@ -74,7 +74,7 @@ python lerobot/scripts/train.py \
 > Notice we added `--output_dir` to explicitly tell where to write outputs from this run (checkpoints, training state, configs etc.). This is not mandatory and if you don't specify it, a default directory will be created from the current date and time, env.type and policy.type. This will typically look like `outputs/train/2025-01-24/16-10-05_aloha_act`.
 
 We now want to train a different policy for aloha on another task. We'll change the dataset and use [lerobot/aloha_sim_transfer_cube_human](https://huggingface.co/datasets/lerobot/aloha_sim_transfer_cube_human) instead. Of course, we also need to change the task of the environment as well to match this other task.
-Looking at the [`AlohaEnv`](../../lerobot/common/envs/configs.py) config, the task is `"AlohaInsertion-v0"` by default, which corresponds to the task we trained on in the command above. The [gym-aloha](https://github.com/huggingface/gym-aloha?tab=readme-ov-file#description) environment also has the `AlohaTransferCube-v0` task which corresponds to this other task we want to train on. Putting this together, we can train this new policy on this different task using:
+Looking at the [`AlohaEnv`](../lerobot/common/envs/configs.py) config, the task is `"AlohaInsertion-v0"` by default, which corresponds to the task we trained on in the command above. The [gym-aloha](https://github.com/huggingface/gym-aloha?tab=readme-ov-file#description) environment also has the `AlohaTransferCube-v0` task which corresponds to this other task we want to train on. Putting this together, we can train this new policy on this different task using:
 ```bash
 python lerobot/scripts/train.py \
     --policy.type=act \

lerobot/common/policies/pi0/configuration_pi0.py

@@ -75,6 +75,10 @@ class PI0Config(PreTrainedConfig):
     train_expert_only: bool = False
     train_state_proj: bool = True
 
+    # Useful if you want to reproduce the training of pi0_base from
+    # the original initialization of PaliGemma (before training on robotics data).
+    paligemma_pretrained_path: str | None = "google/paligemma-3b-pt-224"
+
     # Training presets
     optimizer_lr: float = 2.5e-5
     optimizer_betas: tuple[float, float] = (0.9, 0.95)

lerobot/common/policies/pi0/conversion_scripts/compare_with_jax.py

@@ -66,6 +66,12 @@ def main():
     dataset_meta.stats["observation.state"]["std"] = torch.tensor(
         norm_stats["norm_stats"]["state"]["std"][:num_motors], dtype=torch.float32
     )
+    dataset_meta.stats["action"]["mean"] = torch.tensor(
+        norm_stats["norm_stats"]["actions"]["mean"][:num_motors], dtype=torch.float32
+    )
+    dataset_meta.stats["action"]["std"] = torch.tensor(
+        norm_stats["norm_stats"]["actions"]["std"][:num_motors], dtype=torch.float32
+    )
 
     # Create LeRobot batch from Jax
     batch = {}

lerobot/common/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py

@@ -360,17 +360,20 @@ def convert_pi0_checkpoint(checkpoint_dir: str, precision: str, tokenizer_id: st
             empty_cameras=2,
             adapt_to_pi_aloha=True,
             use_delta_joint_actions_aloha=False,
+            paligemma_pretrained_path=None,
         )
     elif "pi0_aloha_towel" in checkpoint_dir:
         pi0_config = PI0Config(
             adapt_to_pi_aloha=True,
             use_delta_joint_actions_aloha=True,
+            paligemma_pretrained_path=None,
         )
     elif "pi0_base" in checkpoint_dir:
         pi0_config = PI0Config(
             empty_cameras=0,
             adapt_to_pi_aloha=False,
             use_delta_joint_actions_aloha=False,
+            paligemma_pretrained_path=None,
         )
     else:
         raise ValueError()

lerobot/common/policies/pi0/modeling_pi0.py

@@ -44,7 +44,7 @@ python lerobot/scripts/train.py \
 
 Example of using the pi0 pretrained model outside LeRobot training framework:
 ```python
-policy = Pi0Policy.from_pretrained("lerobot/pi0")
+policy = PI0Policy.from_pretrained("lerobot/pi0")
 ```
 
 """
@@ -470,6 +470,7 @@ class PI0FlowMatching(nn.Module):
         paligemma_with_export_config = PaliGemmaWithExpertConfig(
             freeze_vision_encoder=self.config.freeze_vision_encoder,
             train_expert_only=self.config.train_expert_only,
+            paligemma_pretrained_path=self.config.paligemma_pretrained_path,
             attention_implementation=self.config.attention_implementation,
         )
         self.paligemma_with_expert = PaliGemmaWithExpertModel(paligemma_with_export_config)

lerobot/common/policies/pi0/paligemma_with_expert.py

@@ -20,6 +20,7 @@ from pytest import Cache
 from torch import nn
 from transformers import (
     AutoConfig,
+    AutoModel,
     GemmaForCausalLM,
     PaliGemmaForConditionalGeneration,
     PretrainedConfig,
@@ -66,11 +67,13 @@ class PaliGemmaWithExpertConfig(PretrainedConfig):
         gemma_expert_config: dict | None = None,
         freeze_vision_encoder: bool = True,
         train_expert_only: bool = True,
+        paligemma_pretrained_path: str | None = "google/paligemma-3b-pt-224",
         attention_implementation: str = "eager",
         **kwargs,
     ):
         self.freeze_vision_encoder = freeze_vision_encoder
         self.train_expert_only = train_expert_only
+        self.paligemma_pretrained_path = paligemma_pretrained_path
         self.attention_implementation = attention_implementation
 
         if paligemma_config is None:
@@ -167,6 +170,11 @@ class PaliGemmaWithExpertConfig(PretrainedConfig):
                 f"Wrong value provided for `attention_implementation` ({self.attention_implementation}). Expected 'eager', 'fa2' or 'flex'."
             )
 
+        if self.paligemma_pretrained_path is not None and self.paligemma_config is not None:
+            raise ValueError(
+                "When 'paligemma_pretrained_path' is provided, 'paligemma_config' needs to be None."
+            )
+
 
 class PaliGemmaWithExpertModel(PreTrainedModel):
     config_class = PaliGemmaWithExpertConfig
@@ -174,7 +182,10 @@ class PaliGemmaWithExpertModel(PreTrainedModel):
     def __init__(self, config: PaliGemmaWithExpertConfig):
         super().__init__(config=config)
         self.config = config
-        self.paligemma = PaliGemmaForConditionalGeneration(config=config.paligemma_config)
+        if config.paligemma_pretrained_path is not None:
+            self.paligemma = AutoModel(config.paligemma_pretrained_path)
+        else:
+            self.paligemma = PaliGemmaForConditionalGeneration(config=config.paligemma_config)
         self.gemma_expert = GemmaForCausalLM(config=config.gemma_expert_config)
         # Remove unused embed_tokens
         self.gemma_expert.model.embed_tokens = None

lerobot/common/robot_devices/cameras/intelrealsense.py

@@ -512,13 +512,13 @@ if __name__ == "__main__":
     )
     parser.add_argument(
         "--width",
-        type=str,
+        type=int,
         default=640,
         help="Set the width for all cameras. If not provided, use the default width of each camera.",
     )
     parser.add_argument(
         "--height",
-        type=str,
+        type=int,
         default=480,
         help="Set the height for all cameras. If not provided, use the default height of each camera.",
     )

lerobot/common/robot_devices/cameras/opencv.py

@@ -492,13 +492,13 @@ if __name__ == "__main__":
     )
     parser.add_argument(
         "--width",
-        type=str,
+        type=int,
         default=None,
         help="Set the width for all cameras. If not provided, use the default width of each camera.",
     )
     parser.add_argument(
         "--height",
-        type=str,
+        type=int,
         default=None,
         help="Set the height for all cameras. If not provided, use the default height of each camera.",
     )

lerobot/scripts/visualize_dataset_html.py

@@ -174,7 +174,10 @@ def run_server(
                 dataset.meta.get_video_file_path(episode_id, key) for key in dataset.meta.video_keys
             ]
             videos_info = [
-                {"url": url_for("static", filename=video_path), "filename": video_path.parent.name}
+                {
+                    "url": url_for("static", filename=str(video_path).replace("\\", "/")),
+                    "filename": video_path.parent.name,
+                }
                 for video_path in video_paths
             ]
             tasks = dataset.meta.episodes[episode_id]["tasks"]
@@ -381,7 +384,7 @@ def visualize_dataset_html(
         if isinstance(dataset, LeRobotDataset):
             ln_videos_dir = static_dir / "videos"
             if not ln_videos_dir.exists():
-                ln_videos_dir.symlink_to((dataset.root / "videos").resolve())
+                ln_videos_dir.symlink_to((dataset.root / "videos").resolve().as_posix())
 
         if serve:
             run_server(dataset, episodes, host, port, static_dir, template_dir)