Compare commits

..

1 Commits

Author SHA1 Message Date
Khalil Meftah 11dcdc6b23 fix(hub): pin pretrained artifacts to one commit 2026-07-28 14:16:09 +02:00
93 changed files with 886 additions and 2578 deletions
+7 -11
View File
@@ -61,20 +61,16 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
**4.1 Install** **4.1 Install**
```bash ```bash
# uv (recommended — see AGENTS.md and CLAUDE.md) pip install 'lerobot[feetech]' # SO-100/SO-101 motor stack
uv sync --locked --extra feetech # SO-100/SO-101 motor stack # pip install 'lerobot[all]' # everything
# uv sync --locked --extra all # everything # pip install 'lerobot[aloha,pusht]' # specific features
# uv sync --locked --extra smolvla # add SmolVLA deps # pip install 'lerobot[smolvla]' # add SmolVLA deps
# pip (alternative, e.g. when not working from source)
# pip install 'lerobot[feetech]'
# pip install 'lerobot[all]'
# pip install 'lerobot[smolvla]'
git lfs install && git lfs pull git lfs install && git lfs pull
hf auth login # required to push datasets/policies hf auth login # required to push datasets/policies
``` ```
Contributors can alternatively use `uv sync --locked --extra feetech` (see `AGENTS.md`).
**4.2 Find USB ports** — run once per arm, unplug when prompted. **4.2 Find USB ports** — run once per arm, unplug when prompted.
```bash ```bash
+5 -4
View File
@@ -68,16 +68,17 @@ ENV HOME=/home/user_lerobot \
# issues with MuJoCo and OpenGL drivers. # issues with MuJoCo and OpenGL drivers.
RUN uv venv --python python${PYTHON_VERSION} RUN uv venv --python python${PYTHON_VERSION}
# Install third-party dependencies separately for layer caching # Install Python dependencies for caching
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./ COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
RUN uv sync --locked --extra all --no-install-project --no-cache COPY --chown=user_lerobot:user_lerobot src/ src/
RUN uv sync --locked --extra all --no-cache
RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas
# Copy the application source code and install the local project # Copy the rest of the application source code
# Make sure to have the git-LFS files for testing # Make sure to have the git-LFS files for testing
COPY --chown=user_lerobot:user_lerobot . . COPY --chown=user_lerobot:user_lerobot . .
RUN uv sync --locked --extra all --no-cache
# Set the default command # Set the default command
CMD ["/bin/bash"] CMD ["/bin/bash"]
+5 -4
View File
@@ -60,14 +60,15 @@ ENV HOME=/home/user_lerobot \
# run other Python projects in the same container without dependency conflicts. # run other Python projects in the same container without dependency conflicts.
RUN uv venv RUN uv venv
# Install third-party dependencies separately for layer caching # Install Python dependencies for caching
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./ COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
RUN uv sync --locked --extra all --no-install-project --no-cache COPY --chown=user_lerobot:user_lerobot src/ src/
# Copy the application code and install the local project RUN uv sync --locked --extra all --no-cache
# Copy the rest of the application code
# Make sure to have the git-LFS files for testing # Make sure to have the git-LFS files for testing
COPY --chown=user_lerobot:user_lerobot . . COPY --chown=user_lerobot:user_lerobot . .
RUN uv sync --locked --extra all --no-cache
# Set the default command # Set the default command
CMD ["/bin/bash"] CMD ["/bin/bash"]
+3 -3
View File
@@ -58,7 +58,7 @@ final_action = postprocessor(action)
## Hardware API redesign ## Hardware API redesign
PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is an overview of what changed and how you can continue to work with datasets created before this pull request. PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is a overview of what changed and how you can continue to work with datasets created before this pull request.
### What changed? ### What changed?
@@ -129,8 +129,8 @@ python examples/backward_compatibility/replay.py \
Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied. Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied.
To find these transformations, we recommend first replaying an episode of the dataset your policy was trained on using the section above. To find these transformations, we recommend to first try and and replay an episode of the dataset your policy was trained on using the section above.
Then, add these same transformations to your inference script (shown here in the `record.py` script): Then, add these same transformations on your inference script (shown here in the `record.py` script):
```diff ```diff
action_values = predict_action( action_values = predict_action(
-13
View File
@@ -136,10 +136,6 @@ config = RealSenseCameraConfig(
height=480, height=480,
color_mode=ColorMode.RGB, color_mode=ColorMode.RGB,
use_depth=True, use_depth=True,
# Optional fixed color controls. Omit them to leave the current sensor settings unchanged.
exposure=120,
gain=64,
white_balance=4600,
rotation=Cv2Rotation.NO_ROTATION rotation=Cv2Rotation.NO_ROTATION
) )
@@ -158,15 +154,6 @@ finally:
``` ```
<!-- prettier-ignore-end --> <!-- prettier-ignore-end -->
Manual color controls disable the corresponding automatic exposure or white-balance mode. Their
supported ranges vary by camera model; an invalid value raises an error at connection time that
includes the range reported by the sensor. Requesting an unsupported control also raises an error.
Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
require `use_rgb=True`.
On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
exposure or gain also affects the depth stream.
</hfoption> </hfoption>
</hfoptions> </hfoptions>
+15 -2
View File
@@ -88,6 +88,20 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
The same policy can work with different environment processors, and the same environment processor can work with different policies: The same policy can work with different environment processors, and the same environment processor can work with different policies:
````python
# Use SmolVLA policy with LIBERO environment
# Use SmolVLA policy with LIBERO environment
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
env_cfg=libero_cfg,
policy_cfg=smolvla_cfg,
)
smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
# Or use ACT policy with the same LIBERO environment
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
env_cfg=libero_cfg,
policy_cfg=act_cfg,
)
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
```python ```python
# Use SmolVLA policy with LIBERO environment # Use SmolVLA policy with LIBERO environment
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors( libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
@@ -102,7 +116,6 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
policy_cfg=act_cfg, policy_cfg=act_cfg,
) )
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg) act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
```
### 3. **Easier Experimentation** ### 3. **Easier Experimentation**
@@ -132,7 +145,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
state = torch.cat([eef_pos, eef_axisangle, eef_vel, state = torch.cat([eef_pos, eef_axisangle, eef_vel,
gripper_pos, gripper_vel], dim=-1) # 14D gripper_pos, gripper_vel], dim=-1) # 14D
return state return state
``` ````
### 4. **Cleaner Environment Code** ### 4. **Cleaner Environment Code**
+4 -4
View File
@@ -40,10 +40,10 @@ This tutorial guides you through updating the firmware of Feetech motors using t
For each motor you want to update: For each motor you want to update:
1. **Select the motor** from the list by clicking on it 1. **Select the motor** from the list by clicking on it
2. **Click the Upgrade tab**: 2. **Click on Upgrade tab**:
3. **Click the Online button**: 3. **Click on Online button**:
- If a potential firmware update is found, it will be displayed in the box - If an potential firmware update is found, it will be displayed in the box
4. **Click the Upgrade button**: 4. **Click on Upgrade button**:
- The update progress will be displayed - The update progress will be displayed
## Step 6: Verify Update ## Step 6: Verify Update
+1 -1
View File
@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
### Record ### Record
This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data). This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
```bash ```bash
lerobot-record \ lerobot-record \
+1 -1
View File
@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
- [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) for controlling Dynamixel servos - [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) for controlling Dynamixel servos
Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API. Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py) For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py)
Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial): Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
+1 -1
View File
@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
pip install -e ".[lekiwi]" pip install -e ".[lekiwi]"
``` ```
Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖. Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands. Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
# Step-by-Step Assembly Instructions # Step-by-Step Assembly Instructions
+1 -1
View File
@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs
## Reproducing π₀Fast results ## Reproducing π₀Fast results
We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero). We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
The finetuned model can be found here: The finetuned model can be found here:
+4 -4
View File
@@ -22,7 +22,7 @@ With processors, you choose the learning features you want to use for your polic
## Three pipelines ## Three pipelines
We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match. We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
Each of these pipelines handles different conversions between different action and observation spaces. Below is a quick explanation of each pipeline. Each of these pipelines handle different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets) 1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets)
2. Pipeline 2: Dataset action space → robot command space (EE targets → joints) 2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
@@ -74,15 +74,15 @@ In the phone to SO-100 follower examples we use the following adapters:
- `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition. - `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
- `transition_to_robot_action`: transforms the pipeline transition to a robot action dict. - `transition_to_robot_action`: transforms the pipeline transition to a robot action dict.
- `observation_to_transition`: transforms the robot observation dict to a pipeline transition. - `observation_to_transition`: transforms the robot observation dict to a pipeline transition.
- `transition_to_observation`: transforms the pipeline transition to an observation dict. - `transition_to_observation`: transforms the pipeline transition to a observation dict.
Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details. Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
## Dataset feature contracts ## Dataset feature contracts
Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`. Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`.
Below is an example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples: Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
```python ```python
def transform_features( def transform_features(
+2 -2
View File
@@ -57,7 +57,7 @@ policy_cfg.rtc_config = RTCConfig(
policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda") policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")
# Now use predict_action_chunk with RTC parameters # Now use predict_action_chunk with RTC parameters
inference_delay = 4 # How many steps of inference latency, this value should be calculated based on the inference latency of the policy inference_delay = 4 # How many steps of inference latency, this values should be calculated based on the inference latency of the policy
# Initialize the action queue # Initialize the action queue
action_queue = ActionQueue(policy_cfg.rtc_config) action_queue = ActionQueue(policy_cfg.rtc_config)
@@ -100,7 +100,7 @@ Typical values: 8-12 steps
RTCConfig(execution_horizon=10) RTCConfig(execution_horizon=10)
``` ```
**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is an optimal value. **`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is a optimal value.
**`prefix_attention_schedule`**: How to weight consistency across the overlap region. **`prefix_attention_schedule`**: How to weight consistency across the overlap region.
+1 -1
View File
@@ -93,7 +93,7 @@ lerobot-train --help
## Evaluate the finetuned model and run it in real-time ## Evaluate the finetuned model and run it in real-time
Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset). Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
Once you are logged in, you can run inference in your setup by doing: Once you are logged in, you can run inference in your setup by doing:
```bash ```bash
+2 -2
View File
@@ -50,11 +50,11 @@ lerobot-edit-dataset \
Divide a dataset into multiple subsets. Divide a dataset into multiple subsets.
```bash ```bash
# Split by fractions (e.g. 60% train, 20% val, 20% test) # Split by fractions (e.g. 80% train, 20% test, 20% val)
lerobot-edit-dataset \ lerobot-edit-dataset \
--repo_id lerobot/pusht \ --repo_id lerobot/pusht \
--operation.type split \ --operation.type split \
--operation.splits '{"train": 0.6, "val": 0.2, "test": 0.2}' --operation.splits '{"train": 0.8, "test": 0.2, "val": 0.2}'
# Split by specific episode indices # Split by specific episode indices
lerobot-edit-dataset \ lerobot-edit-dataset \
-13
View File
@@ -494,19 +494,6 @@ ignore_errors = true
module = "lerobot.envs.*" module = "lerobot.envs.*"
ignore_errors = false ignore_errors = false
[[tool.mypy.overrides]]
module = "lerobot.annotations.*"
ignore_errors = false
disallow_untyped_defs = true
disallow_incomplete_defs = true
check_untyped_defs = true
[[tool.mypy.overrides]]
module = "lerobot.transforms.*"
ignore_errors = false
disallow_untyped_defs = true
disallow_incomplete_defs = true
check_untyped_defs = true
# [[tool.mypy.overrides]] # [[tool.mypy.overrides]]
# module = "lerobot.utils.*" # module = "lerobot.utils.*"
+48 -78
View File
@@ -120,22 +120,14 @@ class OpenCVCamera(Camera):
self.rotation: int | None = get_cv2_rotation(config.rotation) self.rotation: int | None = get_cv2_rotation(config.rotation)
self.backend: int = config.backend self.backend: int = config.backend
self.capture_width: int | None = None if self.height and self.width:
self.capture_height: int | None = None self.capture_width, self.capture_height = self.width, self.height
self._reset_connection_settings() if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
def __str__(self) -> str: def __str__(self) -> str:
return f"{self.__class__.__name__}({self.index_or_path})" return f"{self.__class__.__name__}({self.index_or_path})"
def _reset_connection_settings(self) -> None:
"""Restore settings that may have been auto-detected during a failed connection."""
self.fps = self.config.fps
self.width = self.config.width
self.height = self.config.height
self.capture_width, self.capture_height = self.width, self.height
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
@property @property
def is_connected(self) -> bool: def is_connected(self) -> bool:
"""Checks if the camera is currently connected and opened.""" """Checks if the camera is currently connected and opened."""
@@ -172,25 +164,17 @@ class OpenCVCamera(Camera):
f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras." f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
) )
try: self._configure_capture_settings()
self._configure_capture_settings() self._start_read_thread()
self._start_read_thread()
if warmup and self.warmup_s > 0: if warmup and self.warmup_s > 0:
start_time = time.time() start_time = time.time()
while time.time() - start_time < self.warmup_s: while time.time() - start_time < self.warmup_s:
self.async_read(timeout_ms=self.warmup_s * 1000) self.async_read(timeout_ms=self.warmup_s * 1000)
time.sleep(0.1) time.sleep(0.1)
with self.frame_lock: with self.frame_lock:
if self.latest_frame is None: if self.latest_frame is None:
raise ConnectionError(f"{self} failed to capture frames during warmup.") raise ConnectionError(f"{self} failed to capture frames during warmup.")
except BaseException:
try:
self._cleanup_resources()
except Exception:
logger.exception(f"Failed to fully clean up {self} after connect() failed.")
self._reset_connection_settings()
raise
logger.info(f"{self} connected.") logger.info(f"{self} connected.")
@@ -328,36 +312,32 @@ class OpenCVCamera(Camera):
for target in targets_to_scan: for target in targets_to_scan:
camera = cv2.VideoCapture(target) camera = cv2.VideoCapture(target)
try: if camera.isOpened():
if camera.isOpened(): default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH)) default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT)) default_fps = camera.get(cv2.CAP_PROP_FPS)
default_fps = camera.get(cv2.CAP_PROP_FPS) default_format = camera.get(cv2.CAP_PROP_FORMAT)
default_format = camera.get(cv2.CAP_PROP_FORMAT)
# Get FOURCC code and convert to string # Get FOURCC code and convert to string
default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC) default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC)
default_fourcc_code_int = int(default_fourcc_code) default_fourcc_code_int = int(default_fourcc_code)
default_fourcc = "".join( default_fourcc = "".join([chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)])
[chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)]
)
camera_info = { camera_info = {
"name": f"OpenCV Camera @ {target}", "name": f"OpenCV Camera @ {target}",
"type": "OpenCV", "type": "OpenCV",
"id": target, "id": target,
"backend_api": camera.getBackendName(), "backend_api": camera.getBackendName(),
"default_stream_profile": { "default_stream_profile": {
"format": default_format, "format": default_format,
"fourcc": default_fourcc, "fourcc": default_fourcc,
"width": default_width, "width": default_width,
"height": default_height, "height": default_height,
"fps": default_fps, "fps": default_fps,
}, },
} }
found_cameras_info.append(camera_info) found_cameras_info.append(camera_info)
finally:
camera.release() camera.release()
return found_cameras_info return found_cameras_info
@@ -516,26 +496,6 @@ class OpenCVCamera(Camera):
self.latest_timestamp = None self.latest_timestamp = None
self.new_frame_event.clear() self.new_frame_event.clear()
def _cleanup_resources(self) -> None:
"""Stop background reads and release the capture, including after partial setup."""
read_thread = self.thread
videocapture = self.videocapture
try:
self._stop_read_thread()
finally:
self.videocapture = None
try:
if videocapture is not None:
videocapture.release()
finally:
# Releasing the device may unblock a hardware read that outlived
# the first bounded join in _stop_read_thread().
if read_thread is not None and read_thread.is_alive():
read_thread.join(timeout=2.0)
if read_thread.is_alive(): # pragma: no cover
logger.warning(f"{self} read thread remained alive after releasing the capture.")
@check_if_not_connected @check_if_not_connected
def async_read(self, timeout_ms: float = 200) -> NDArray[Any]: def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
""" """
@@ -626,6 +586,16 @@ class OpenCVCamera(Camera):
if not self.is_connected and self.thread is None: if not self.is_connected and self.thread is None:
raise DeviceNotConnectedError(f"{self} not connected.") raise DeviceNotConnectedError(f"{self} not connected.")
self._cleanup_resources() if self.thread is not None:
self._stop_read_thread()
if self.videocapture is not None:
self.videocapture.release()
self.videocapture = None
with self.frame_lock:
self.latest_frame = None
self.latest_timestamp = None
self.new_frame_event.clear()
logger.info(f"{self} disconnected.") logger.info(f"{self} disconnected.")
+33 -171
View File
@@ -121,9 +121,6 @@ class RealSenseCamera(Camera):
self.config = config self.config = config
self.width: int | None = config.width
self.height: int | None = config.height
if config.serial_number_or_name.isdigit(): if config.serial_number_or_name.isdigit():
self.serial_number = config.serial_number_or_name self.serial_number = config.serial_number_or_name
else: else:
@@ -134,9 +131,6 @@ class RealSenseCamera(Camera):
self.use_rgb = config.use_rgb self.use_rgb = config.use_rgb
self.use_depth = config.use_depth self.use_depth = config.use_depth
self.warmup_s = config.warmup_s self.warmup_s = config.warmup_s
self.exposure: int | None = config.exposure
self.gain: int | None = config.gain
self.white_balance: int | None = config.white_balance
self.rs_pipeline: rs.pipeline | None = None self.rs_pipeline: rs.pipeline | None = None
self.rs_profile: rs.pipeline_profile | None = None self.rs_profile: rs.pipeline_profile | None = None
@@ -151,23 +145,14 @@ class RealSenseCamera(Camera):
self.rotation: int | None = get_cv2_rotation(config.rotation) self.rotation: int | None = get_cv2_rotation(config.rotation)
self.capture_width: int | None = None if self.height and self.width:
self.capture_height: int | None = None self.capture_width, self.capture_height = self.width, self.height
self._reset_connection_settings() if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
def __str__(self) -> str: def __str__(self) -> str:
return f"{self.__class__.__name__}({self.serial_number})" return f"{self.__class__.__name__}({self.serial_number})"
def _reset_connection_settings(self) -> None:
"""Restore settings that may have been auto-detected during a failed connection."""
self.fps = self.config.fps
self.width = self.config.width
self.height = self.config.height
self.warmup_s = self.config.warmup_s
self.capture_width, self.capture_height = self.width, self.height
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
self.capture_width, self.capture_height = self.height, self.width
@property @property
def is_connected(self) -> bool: def is_connected(self) -> bool:
"""Checks if the camera pipeline is started and streams are active.""" """Checks if the camera pipeline is started and streams are active."""
@@ -187,8 +172,7 @@ class RealSenseCamera(Camera):
Raises: Raises:
DeviceAlreadyConnectedError: If the camera is already connected. DeviceAlreadyConnectedError: If the camera is already connected.
ValueError: If the configuration is invalid, a requested sensor option is unsupported, ValueError: If the configuration is invalid (e.g., missing serial/name, name not unique).
or a requested sensor value is invalid.
ConnectionError: If the camera is found but fails to start the pipeline or no RealSense devices are detected at all. ConnectionError: If the camera is found but fails to start the pipeline or no RealSense devices are detected at all.
RuntimeError: If the pipeline starts but fails to apply requested settings. RuntimeError: If the pipeline starts but fails to apply requested settings.
""" """
@@ -206,31 +190,22 @@ class RealSenseCamera(Camera):
f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras." f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
) from e ) from e
try: self._configure_capture_settings()
self._configure_capture_settings() self._start_read_thread()
self._configure_sensor_options()
self._start_read_thread()
# NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise. # NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise.
self.warmup_s = max(self.warmup_s, 1) self.warmup_s = max(self.warmup_s, 1)
warmup_read = self.async_read if self.use_rgb else self.async_read_depth warmup_read = self.async_read if self.use_rgb else self.async_read_depth
start_time = time.time() start_time = time.time()
while time.time() - start_time < self.warmup_s: while time.time() - start_time < self.warmup_s:
warmup_read(timeout_ms=self.warmup_s * 1000) warmup_read(timeout_ms=self.warmup_s * 1000)
time.sleep(0.1) time.sleep(0.1)
with self.frame_lock: with self.frame_lock:
if (self.use_rgb and self.latest_color_frame is None) or ( if (self.use_rgb and self.latest_color_frame is None) or (
self.use_depth and self.latest_depth_frame is None self.use_depth and self.latest_depth_frame is None
): ):
raise ConnectionError(f"{self} failed to capture frames during warmup.") raise ConnectionError(f"{self} failed to capture frames during warmup.")
except BaseException:
try:
self._cleanup_resources()
except Exception:
logger.exception(f"Failed to fully clean up {self} after connect() failed.")
self._reset_connection_settings()
raise
logger.info(f"{self} connected.") logger.info(f"{self} connected.")
@@ -364,111 +339,6 @@ class RealSenseCamera(Camera):
self.new_frame_event.clear() self.new_frame_event.clear()
return self._async_read(timeout_ms=10000, read_depth=read_depth) return self._async_read(timeout_ms=10000, read_depth=read_depth)
def _get_color_sensor(self) -> "rs.sensor":
"""Returns the sensor that controls the color stream.
Most RealSense cameras expose "RGB Camera" for color. The D405 has no
separate RGB module its color stream comes from "Stereo Module".
We try RGB Camera first, then fall back to Stereo Module.
"""
if self.rs_profile is None:
raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
device = self.rs_profile.get_device()
sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
for name in ("RGB Camera", "Stereo Module"):
if name in sensors:
return sensors[name]
available = list(sensors.keys())
raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
"""Sets a sensor option, re-raising range errors with actionable diagnostics."""
try:
sensor.set_option(option, value)
except Exception as e:
range_info = ""
try:
option_range = sensor.get_option_range(option)
range_info = (
f" (supported range: min={option_range.min}, max={option_range.max}, "
f"step={option_range.step}, default={option_range.default})"
)
except Exception:
range_info = " (option range unavailable)"
raise ValueError(
f"{self}: failed to set {label} to {value}{range_info}. Original error: {e}"
) from e
def _configure_sensor_options(self) -> None:
"""Applies manual sensor options (exposure, gain, white balance) to the color sensor.
When exposure or gain is set, auto-exposure is disabled first. When white_balance
is set, auto white balance is disabled first. An omitted option is left unchanged,
and configuration is skipped entirely if all options are omitted.
Raises:
ValueError: If the sensor does not support a requested option or a requested
value is invalid. Invalid-value errors include the option name, requested
value, and supported range when available.
"""
if self.exposure is None and self.gain is None and self.white_balance is None:
return
color_sensor = self._get_color_sensor()
requested_options = (
(rs.option.exposure, self.exposure, "exposure"),
(rs.option.gain, self.gain, "gain"),
(rs.option.white_balance, self.white_balance, "white balance"),
)
unsupported_options = [
label
for option, value, label in requested_options
if value is not None and not color_sensor.supports(option)
]
if unsupported_options:
raise ValueError(
f"{self}: color sensor does not support requested manual options: {unsupported_options}."
)
manual_exposure_requested = self.exposure is not None or self.gain is not None
if manual_exposure_requested:
if color_sensor.supports(rs.option.enable_auto_exposure):
self._set_sensor_option(color_sensor, rs.option.enable_auto_exposure, 0, "auto-exposure")
logger.info(f"{self} auto-exposure disabled.")
else:
logger.warning(
f"{self} sensor does not support disabling auto-exposure; "
"applying manual exposure/gain directly."
)
if self.exposure is not None:
self._set_sensor_option(color_sensor, rs.option.exposure, self.exposure, "exposure")
logger.info(f"{self} exposure set to {self.exposure}.")
if self.gain is not None:
self._set_sensor_option(color_sensor, rs.option.gain, self.gain, "gain")
logger.info(f"{self} gain set to {self.gain}.")
if self.white_balance is not None:
if color_sensor.supports(rs.option.enable_auto_white_balance):
self._set_sensor_option(
color_sensor, rs.option.enable_auto_white_balance, 0, "auto white balance"
)
logger.info(f"{self} auto white balance disabled.")
else:
logger.warning(
f"{self} sensor does not support disabling auto white balance; "
"applying manual white balance directly."
)
self._set_sensor_option(
color_sensor, rs.option.white_balance, self.white_balance, "white balance"
)
logger.info(f"{self} white balance set to {self.white_balance}.")
@check_if_not_connected @check_if_not_connected
def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]: def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
""" """
@@ -671,27 +541,6 @@ class RealSenseCamera(Camera):
self.latest_timestamp = None self.latest_timestamp = None
self.new_frame_event.clear() self.new_frame_event.clear()
def _cleanup_resources(self) -> None:
"""Stop background reads and stop the pipeline, including after partial setup."""
read_thread = self.thread
rs_pipeline = self.rs_pipeline
try:
self._stop_read_thread()
finally:
self.rs_pipeline = None
self.rs_profile = None
try:
if rs_pipeline is not None:
rs_pipeline.stop()
finally:
# Stopping the pipeline may unblock a hardware read that outlived
# the first bounded join in _stop_read_thread().
if read_thread is not None and read_thread.is_alive():
read_thread.join(timeout=2.0)
if read_thread.is_alive(): # pragma: no cover
logger.warning(f"{self} read thread remained alive after stopping the pipeline.")
def _async_read(self, timeout_ms: float, read_depth: bool = False) -> NDArray[Any]: def _async_read(self, timeout_ms: float, read_depth: bool = False) -> NDArray[Any]:
"""Shared helper for :meth:`async_read`/:meth:`async_read_depth`: return the latest buffered frame.""" """Shared helper for :meth:`async_read`/:meth:`async_read_depth`: return the latest buffered frame."""
if self.thread is None or not self.thread.is_alive(): if self.thread is None or not self.thread.is_alive():
@@ -835,5 +684,18 @@ class RealSenseCamera(Camera):
f"Attempted to disconnect {self}, but it appears already disconnected." f"Attempted to disconnect {self}, but it appears already disconnected."
) )
self._cleanup_resources() if self.thread is not None:
self._stop_read_thread()
if self.rs_pipeline is not None:
self.rs_pipeline.stop()
self.rs_pipeline = None
self.rs_profile = None
with self.frame_lock:
self.latest_color_frame = None
self.latest_depth_frame = None
self.latest_timestamp = None
self.new_frame_event.clear()
logger.info(f"{self} disconnected.") logger.info(f"{self} disconnected.")
@@ -46,17 +46,6 @@ class RealSenseCameraConfig(CameraConfig):
use_depth: Whether to enable depth stream. Defaults to False. use_depth: Whether to enable depth stream. Defaults to False.
rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation. rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
warmup_s: Time reading frames before returning from connect (in seconds) warmup_s: Time reading frames before returning from connect (in seconds)
exposure: Manual exposure value for the color sensor. When set, auto-exposure is
disabled and this fixed value is used. Valid ranges are camera-model specific
and reported if the value is rejected. Defaults to None (leave unchanged).
gain: Manual gain value for the color sensor. When set, auto-exposure is disabled
and this fixed gain is used, which also freezes exposure at its current value
when no exposure is configured. Valid ranges are camera-model specific and
reported if the value is rejected. Defaults to None (leave unchanged).
white_balance: Manual white balance value for the color sensor. When set, auto
white balance is disabled and this fixed value is used. Valid ranges are
camera-model specific and reported if the value is rejected. Defaults to None
(leave unchanged).
Note: Note:
- Either name or serial_number must be specified. - Either name or serial_number must be specified.
@@ -72,9 +61,6 @@ class RealSenseCameraConfig(CameraConfig):
use_depth: bool = False use_depth: bool = False
rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
warmup_s: int = 1 warmup_s: int = 1
exposure: int | None = None
gain: int | None = None
white_balance: int | None = None
def __post_init__(self) -> None: def __post_init__(self) -> None:
self.color_mode = ColorMode(self.color_mode) self.color_mode = ColorMode(self.color_mode)
@@ -83,18 +69,6 @@ class RealSenseCameraConfig(CameraConfig):
if not self.use_rgb and not self.use_depth: if not self.use_rgb and not self.use_depth:
raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.") raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.")
manual_color_options = {
"exposure": self.exposure,
"gain": self.gain,
"white_balance": self.white_balance,
}
configured_color_options = [name for name, value in manual_color_options.items() if value is not None]
if configured_color_options and not self.use_rgb:
raise ValueError(
"Manual color sensor options require `use_rgb=True`. "
f"Configured options: {configured_color_options}."
)
values = (self.fps, self.width, self.height) values = (self.fps, self.width, self.height)
if any(v is not None for v in values) and any(v is None for v in values): if any(v is not None for v in values) and any(v is None for v in values):
raise ValueError( raise ValueError(
-6
View File
@@ -71,19 +71,13 @@ class DatasetRecordConfig:
# Number of threads per encoder instance. None = auto (codec default). # Number of threads per encoder instance. None = auto (codec default).
# Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc.. # Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
encoder_threads: int | None = None encoder_threads: int | None = None
# Skip appending the date-time tag to repo_id, keeping the user-provided name as-is
# (e.g. self-managed versioned names intended for a later `lerobot-edit-dataset merge`).
no_stamp: bool = False
def stamp_repo_id(self) -> None: def stamp_repo_id(self) -> None:
"""Append a date-time tag to ``repo_id`` so each recording session gets a unique name. """Append a date-time tag to ``repo_id`` so each recording session gets a unique name.
Must be called explicitly at dataset *creation* time not on resume, Must be called explicitly at dataset *creation* time not on resume,
where the existing ``repo_id`` (already stamped) must be preserved. where the existing ``repo_id`` (already stamped) must be preserved.
No-op when ``no_stamp`` is set, preserving a user-managed ``repo_id``.
""" """
if self.no_stamp:
return
if self.repo_id: if self.repo_id:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
self.repo_id = f"{self.repo_id}_{timestamp}" self.repo_id = f"{self.repo_id}_{timestamp}"
+6 -1
View File
@@ -48,8 +48,13 @@ class EvalPipelineConfig:
if policy_path: if policy_path:
yaml_overrides = parser.get_yaml_overrides("policy") yaml_overrides = parser.get_yaml_overrides("policy")
cli_overrides = parser.get_cli_overrides("policy") or [] cli_overrides = parser.get_cli_overrides("policy") or []
pretrained_revision = parser.parse_arg("pretrained_revision", cli_overrides)
if pretrained_revision is None:
pretrained_revision = parser.parse_arg("pretrained_revision", yaml_overrides)
self.policy = PreTrainedConfig.from_pretrained( self.policy = PreTrainedConfig.from_pretrained(
policy_path, cli_overrides=yaml_overrides + cli_overrides policy_path,
revision=pretrained_revision,
cli_overrides=yaml_overrides + cli_overrides,
) )
self.policy.pretrained_path = Path(policy_path) self.policy.pretrained_path = Path(policy_path)
+30 -3
View File
@@ -29,7 +29,7 @@ from huggingface_hub.errors import HfHubHTTPError
from lerobot.optim import LRSchedulerConfig, OptimizerConfig from lerobot.optim import LRSchedulerConfig, OptimizerConfig
from lerobot.utils.constants import ACTION, OBS_STATE from lerobot.utils.constants import ACTION, OBS_STATE
from lerobot.utils.device_utils import auto_select_torch_device, is_amp_available, is_torch_device_available from lerobot.utils.device_utils import auto_select_torch_device, is_amp_available, is_torch_device_available
from lerobot.utils.hub import HubMixin from lerobot.utils.hub import HubMixin, extract_commit_hash
from .types import FeatureType, PolicyFeature from .types import FeatureType, PolicyFeature
@@ -82,6 +82,26 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
# Optional Hub revision (commit hash, branch, or tag) to pin the pretrained model version. # Optional Hub revision (commit hash, branch, or tag) to pin the pretrained model version.
pretrained_revision: str | None = None pretrained_revision: str | None = None
@property
def _commit_hash(self) -> str | None:
"""Resolved Hub commit for this runtime load; never serialized."""
return self.__dict__.get("_runtime_commit_hash")
@property
def _commit_hash_source(self) -> str | None:
"""Hub repo whose revision resolved to ``_commit_hash``."""
return self.__dict__.get("_runtime_commit_hash_source")
def _set_hub_commit_hash(self, commit_hash: str | None, source: str | None) -> None:
self.__dict__["_runtime_commit_hash"] = commit_hash
self.__dict__["_runtime_commit_hash_source"] = source if commit_hash is not None else None
def get_hub_revision(self, source: str | Path | None, revision: str | None = None) -> str | None:
"""Return the pinned revision when ``source`` owns the resolved commit."""
if self._commit_hash is not None and self._commit_hash_source == str(source):
return self._commit_hash
return revision
def __post_init__(self) -> None: def __post_init__(self) -> None:
if not self.device or not is_torch_device_available(self.device): if not self.device or not is_torch_device_available(self.device):
auto_device = auto_select_torch_device() auto_device = auto_select_torch_device()
@@ -182,7 +202,8 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
) -> T: ) -> T:
model_id = str(pretrained_name_or_path) model_id = str(pretrained_name_or_path)
config_file: str | None = None config_file: str | None = None
if Path(model_id).is_dir(): is_local = Path(model_id).is_dir()
if is_local:
if CONFIG_NAME in os.listdir(model_id): if CONFIG_NAME in os.listdir(model_id):
config_file = os.path.join(model_id, CONFIG_NAME) config_file = os.path.join(model_id, CONFIG_NAME)
else: else:
@@ -208,8 +229,12 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
if config_file is None: if config_file is None:
raise FileNotFoundError(f"{CONFIG_NAME} not found in {model_id}") raise FileNotFoundError(f"{CONFIG_NAME} not found in {model_id}")
commit_hash = None if is_local else extract_commit_hash(config_file, revision)
with open(config_file) as f: with open(config_file) as f:
config = json.load(f) config = json.load(f)
# Runtime Hub metadata must never become part of the serialized config schema.
config.pop("_commit_hash", None)
config.pop("_commit_hash_source", None)
# Resolve the concrete config subclass from the serialized "type" tag, then parse # Resolve the concrete config subclass from the serialized "type" tag, then parse
# the config (with CLI overrides) directly for that class. The "type" key is # the config (with CLI overrides) directly for that class. The "type" key is
@@ -231,4 +256,6 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
cli_overrides = policy_kwargs.pop("cli_overrides", []) cli_overrides = policy_kwargs.pop("cli_overrides", [])
with draccus.config_type("json"): with draccus.config_type("json"):
return draccus.parse(config_cls, config_file, args=cli_overrides) parsed_config = draccus.parse(config_cls, config_file, args=cli_overrides)
parsed_config._set_hub_commit_hash(commit_hash, model_id)
return parsed_config
+10 -2
View File
@@ -172,8 +172,16 @@ class TrainPipelineConfig(HubMixin):
) )
self.reward_model.pretrained_path = str(Path(reward_model_path)) self.reward_model.pretrained_path = str(Path(reward_model_path))
elif policy_path: elif policy_path:
overrides = parser.get_yaml_overrides("policy") + (parser.get_cli_overrides("policy") or []) yaml_overrides = parser.get_yaml_overrides("policy")
self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=overrides) cli_overrides = parser.get_cli_overrides("policy") or []
pretrained_revision = parser.parse_arg("pretrained_revision", cli_overrides)
if pretrained_revision is None:
pretrained_revision = parser.parse_arg("pretrained_revision", yaml_overrides)
self.policy = PreTrainedConfig.from_pretrained(
policy_path,
revision=pretrained_revision,
cli_overrides=yaml_overrides + cli_overrides,
)
self.policy.pretrained_path = Path(policy_path) self.policy.pretrained_path = Path(policy_path)
elif self.resume: elif self.resume:
self._resolve_resume_checkpoint() self._resolve_resume_checkpoint()
+58 -114
View File
@@ -19,7 +19,6 @@ import copy
import logging import logging
import shutil import shutil
from pathlib import Path from pathlib import Path
from typing import Any, NotRequired, TypedDict
import datasets import datasets
import pandas as pd import pandas as pd
@@ -50,32 +49,8 @@ from .utils import (
) )
from .video_utils import concatenate_video_files, get_video_duration_in_s from .video_utils import concatenate_video_files, get_video_duration_in_s
logger = logging.getLogger(__name__)
type FeatureDict = dict[str, dict[str, Any]] def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> dict[str, dict]:
type ChunkFile = tuple[int, int]
class IndexState(TypedDict):
chunk: int
file: int
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
class VideoIndex(TypedDict):
chunk: int
file: int
latest_duration: float
episode_duration: float
src_to_offset: NotRequired[dict[ChunkFile, float]]
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
dst_file_durations: NotRequired[dict[ChunkFile, float]]
type VideoIndexState = dict[str, VideoIndex]
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> FeatureDict:
"""Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged. """Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged.
Args: Args:
@@ -84,14 +59,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
Returns: Returns:
dict: A dictionary of merged video feature info. dict: A dictionary of merged video feature info.
""" """
merged_info: FeatureDict = copy.deepcopy(all_metadata[0].features) merged_info = copy.deepcopy(all_metadata[0].features)
video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"] video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"]
for vk in video_keys: for vk in video_keys:
video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata] video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
base_video_info = video_infos[0] base_video_info = video_infos[0]
merged_encoder_info: dict[str, Any] = {} merged_encoder_info: dict = {}
fallback_keys: list[str] = [] fallback_keys: list[str] = []
for info_key in VIDEO_ENCODER_INFO_KEYS: for info_key in VIDEO_ENCODER_INFO_KEYS:
values = [info.get(info_key, None) for info in video_infos] values = [info.get(info_key, None) for info in video_infos]
@@ -105,7 +80,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None
if fallback_keys: if fallback_keys:
logger.warning( logging.warning(
f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. " f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
f"Setting these keys to null: {fallback_keys}.", f"Setting these keys to null: {fallback_keys}.",
) )
@@ -117,7 +92,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
return merged_info return merged_info
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[int, str | None, FeatureDict]: def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
"""Validates that all dataset metadata have consistent properties. """Validates that all dataset metadata have consistent properties.
Ensures all datasets have the same fps, robot_type, and features to guarantee Ensures all datasets have the same fps, robot_type, and features to guarantee
@@ -154,9 +129,7 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[i
return fps, robot_type, features return fps, robot_type, features
def update_data_df( def update_data_df(df, src_meta, dst_meta):
df: pd.DataFrame, src_meta: LeRobotDatasetMetadata, dst_meta: LeRobotDatasetMetadata
) -> pd.DataFrame:
"""Updates a data DataFrame with new indices and task mappings for aggregation. """Updates a data DataFrame with new indices and task mappings for aggregation.
Adjusts episode indices, frame indices, and task indices to account for Adjusts episode indices, frame indices, and task indices to account for
@@ -181,12 +154,12 @@ def update_data_df(
def update_meta_data( def update_meta_data(
df: pd.DataFrame, df,
dst_meta: LeRobotDatasetMetadata, dst_meta,
meta_idx: IndexState, meta_idx,
data_idx: IndexState, data_idx,
videos_idx: VideoIndexState, videos_idx,
) -> pd.DataFrame: ):
"""Updates metadata DataFrame with new chunk, file, and timestamp indices. """Updates metadata DataFrame with new chunk, file, and timestamp indices.
Adjusts all indices and timestamps to account for previously aggregated Adjusts all indices and timestamps to account for previously aggregated
@@ -316,7 +289,7 @@ def aggregate_datasets(
chunk_size: int | None = None, chunk_size: int | None = None,
concatenate_videos: bool = True, concatenate_videos: bool = True,
concatenate_data: bool = True, concatenate_data: bool = True,
) -> None: ):
"""Aggregates multiple LeRobot datasets into a single unified dataset. """Aggregates multiple LeRobot datasets into a single unified dataset.
This is the main function that orchestrates the aggregation process by: This is the main function that orchestrates the aggregation process by:
@@ -336,7 +309,7 @@ def aggregate_datasets(
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards. concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
concatenate_data: When False, keep one parquet per source file instead of packing into shards. concatenate_data: When False, keep one parquet per source file instead of packing into shards.
""" """
logger.info("Start aggregate_datasets") logging.info("Start aggregate_datasets")
if data_files_size_in_mb is None: if data_files_size_in_mb is None:
data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
@@ -368,15 +341,15 @@ def aggregate_datasets(
video_files_size_in_mb=video_files_size_in_mb, video_files_size_in_mb=video_files_size_in_mb,
) )
logger.info("Find all tasks") logging.info("Find all tasks")
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique() unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
dst_meta.tasks = pd.DataFrame( dst_meta.tasks = pd.DataFrame(
{"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task") {"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
) )
meta_idx: IndexState = {"chunk": 0, "file": 0} meta_idx = {"chunk": 0, "file": 0}
data_idx: IndexState = {"chunk": 0, "file": 0} data_idx = {"chunk": 0, "file": 0}
videos_idx: VideoIndexState = { videos_idx = {
key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
} }
@@ -400,17 +373,12 @@ def aggregate_datasets(
dst_meta.info.total_frames += src_meta.total_frames dst_meta.info.total_frames += src_meta.total_frames
finalize_aggregation(dst_meta, all_metadata) finalize_aggregation(dst_meta, all_metadata)
logger.info("Aggregation complete.") logging.info("Aggregation complete.")
def aggregate_videos( def aggregate_videos(
src_meta: LeRobotDatasetMetadata, src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True
dst_meta: LeRobotDatasetMetadata, ):
videos_idx: VideoIndexState,
video_files_size_in_mb: float,
chunk_size: int,
concatenate_videos: bool = True,
) -> VideoIndexState:
"""Aggregates video chunks from a source dataset into the destination dataset. """Aggregates video chunks from a source dataset into the destination dataset.
Handles video file concatenation and rotation based on file size limits. Handles video file concatenation and rotation based on file size limits.
@@ -438,16 +406,15 @@ def aggregate_videos(
videos_idx[key]["dst_file_durations"] = {} videos_idx[key]["dst_file_durations"] = {}
for key, video_idx in videos_idx.items(): for key, video_idx in videos_idx.items():
unique_chunk_file_pairs: list[ChunkFile] = sorted( unique_chunk_file_pairs = {
{ (chunk, file)
(chunk, file) for chunk, file in zip(
for chunk, file in zip( src_meta.episodes[f"videos/{key}/chunk_index"],
src_meta.episodes[f"videos/{key}/chunk_index"], src_meta.episodes[f"videos/{key}/file_index"],
src_meta.episodes[f"videos/{key}/file_index"], strict=False,
strict=False, )
) }
} unique_chunk_file_pairs = sorted(unique_chunk_file_pairs)
)
chunk_idx = video_idx["chunk"] chunk_idx = video_idx["chunk"]
file_idx = video_idx["file"] file_idx = video_idx["file"]
@@ -522,14 +489,7 @@ def aggregate_videos(
return videos_idx return videos_idx
def aggregate_data( def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data=True):
src_meta: LeRobotDatasetMetadata,
dst_meta: LeRobotDatasetMetadata,
data_idx: IndexState,
data_files_size_in_mb: float,
chunk_size: int,
concatenate_data: bool = True,
) -> IndexState:
"""Aggregates data chunks from a source dataset into the destination dataset. """Aggregates data chunks from a source dataset into the destination dataset.
Reads source data files, updates indices to match the aggregated dataset, Reads source data files, updates indices to match the aggregated dataset,
@@ -550,16 +510,14 @@ def aggregate_data(
Returns: Returns:
dict: Updated data_idx with current chunk and file indices. dict: Updated data_idx with current chunk and file indices.
""" """
unique_chunk_file_ids: list[ChunkFile] = sorted( unique_chunk_file_ids = {
{ (c, f)
(c, f) for c, f in zip(
for c, f in zip( src_meta.episodes["data/chunk_index"], src_meta.episodes["data/file_index"], strict=False
src_meta.episodes["data/chunk_index"], )
src_meta.episodes["data/file_index"], }
strict=False,
) unique_chunk_file_ids = sorted(unique_chunk_file_ids)
}
)
contains_images = len(dst_meta.image_keys) > 0 contains_images = len(dst_meta.image_keys) > 0
# retrieve features schema for proper image typing in parquet # retrieve features schema for proper image typing in parquet
@@ -567,7 +525,7 @@ def aggregate_data(
# Track source to destination file mapping for metadata update # Track source to destination file mapping for metadata update
# This is critical for handling datasets that are already results of a merge # This is critical for handling datasets that are already results of a merge
src_to_dst: dict[ChunkFile, ChunkFile] = {} src_to_dst: dict[tuple[int, int], tuple[int, int]] = {}
for src_chunk_idx, src_file_idx in unique_chunk_file_ids: for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
src_path = src_meta.root / DEFAULT_DATA_PATH.format( src_path = src_meta.root / DEFAULT_DATA_PATH.format(
@@ -606,13 +564,7 @@ def aggregate_data(
return data_idx return data_idx
def aggregate_metadata( def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
src_meta: LeRobotDatasetMetadata,
dst_meta: LeRobotDatasetMetadata,
meta_idx: IndexState,
data_idx: IndexState,
videos_idx: VideoIndexState,
) -> IndexState:
"""Aggregates metadata from a source dataset into the destination dataset. """Aggregates metadata from a source dataset into the destination dataset.
Reads source metadata files, updates all indices and timestamps, Reads source metadata files, updates all indices and timestamps,
@@ -628,16 +580,16 @@ def aggregate_metadata(
Returns: Returns:
dict: Updated meta_idx with current chunk and file indices. dict: Updated meta_idx with current chunk and file indices.
""" """
chunk_file_ids: list[ChunkFile] = sorted( chunk_file_ids = {
{ (c, f)
(c, f) for c, f in zip(
for c, f in zip( src_meta.episodes["meta/episodes/chunk_index"],
src_meta.episodes["meta/episodes/chunk_index"], src_meta.episodes["meta/episodes/file_index"],
src_meta.episodes["meta/episodes/file_index"], strict=False,
strict=False, )
) }
}
) chunk_file_ids = sorted(chunk_file_ids)
for chunk_idx, file_idx in chunk_file_ids: for chunk_idx, file_idx in chunk_file_ids:
src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx) src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
df = pd.read_parquet(src_path) df = pd.read_parquet(src_path)
@@ -670,16 +622,16 @@ def aggregate_metadata(
def append_or_create_parquet_file( def append_or_create_parquet_file(
df: pd.DataFrame, df: pd.DataFrame,
src_path: Path, src_path: Path,
idx: IndexState, idx: dict[str, int],
max_mb: float, max_mb: float,
chunk_size: int, chunk_size: int,
default_path: str, default_path: str,
contains_images: bool = False, contains_images: bool = False,
aggr_root: Path | None = None, aggr_root: Path = None,
hf_features: datasets.Features | None = None, hf_features: datasets.Features | None = None,
concatenate: bool = True, concatenate: bool = True,
one_row_group_per_episode: bool = False, one_row_group_per_episode: bool = False,
) -> tuple[IndexState, ChunkFile]: ) -> tuple[dict[str, int], tuple[int, int]]:
"""Appends data to an existing parquet file or creates a new one based on size constraints. """Appends data to an existing parquet file or creates a new one based on size constraints.
Manages file rotation when size limits are exceeded to prevent individual files Manages file rotation when size limits are exceeded to prevent individual files
@@ -702,13 +654,7 @@ def append_or_create_parquet_file(
Returns: Returns:
tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
and (dst_chunk, dst_file) is the actual destination file the data was written to. and (dst_chunk, dst_file) is the actual destination file the data was written to.
Raises:
ValueError: If aggr_root is not provided.
""" """
if aggr_root is None:
raise ValueError("aggr_root must be provided.")
dst_chunk, dst_file = idx["chunk"], idx["file"] dst_chunk, dst_file = idx["chunk"], idx["file"]
dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file) dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
@@ -752,9 +698,7 @@ def append_or_create_parquet_file(
return idx, (dst_chunk, dst_file) return idx, (dst_chunk, dst_file)
def finalize_aggregation( def finalize_aggregation(aggr_meta, all_metadata):
aggr_meta: LeRobotDatasetMetadata, all_metadata: list[LeRobotDatasetMetadata]
) -> None:
"""Finalizes the dataset aggregation by writing summary files and statistics. """Finalizes the dataset aggregation by writing summary files and statistics.
Writes the tasks file, info file with total counts and splits, and Writes the tasks file, info file with total counts and splits, and
@@ -764,16 +708,16 @@ def finalize_aggregation(
aggr_meta: Aggregated dataset metadata. aggr_meta: Aggregated dataset metadata.
all_metadata: List of all source dataset metadata objects. all_metadata: List of all source dataset metadata objects.
""" """
logger.info("write tasks") logging.info("write tasks")
write_tasks(aggr_meta.tasks, aggr_meta.root) write_tasks(aggr_meta.tasks, aggr_meta.root)
logger.info("write info") logging.info("write info")
aggr_meta.info.total_tasks = len(aggr_meta.tasks) aggr_meta.info.total_tasks = len(aggr_meta.tasks)
aggr_meta.info.total_episodes = sum(m.total_episodes for m in all_metadata) aggr_meta.info.total_episodes = sum(m.total_episodes for m in all_metadata)
aggr_meta.info.total_frames = sum(m.total_frames for m in all_metadata) aggr_meta.info.total_frames = sum(m.total_frames for m in all_metadata)
aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"} aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"}
write_info(aggr_meta.info, aggr_meta.root) write_info(aggr_meta.info, aggr_meta.root)
logger.info("write stats") logging.info("write stats")
aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata]) aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
write_stats(aggr_meta.stats, aggr_meta.root) write_stats(aggr_meta.stats, aggr_meta.root)
+2 -2
View File
@@ -188,8 +188,8 @@ class LeRobotDatasetMetadata:
def _load_metadata(self): def _load_metadata(self):
self.info = load_info(self.root) self.info = load_info(self.root)
check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION) check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
self.tasks = load_tasks(self.root) if self.total_tasks > 0 else None self.tasks = load_tasks(self.root)
self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None self.episodes = load_episodes(self.root)
self.stats = load_stats(self.root) self.stats = load_stats(self.root)
def ensure_readable(self) -> None: def ensure_readable(self) -> None:
+1 -6
View File
@@ -384,12 +384,7 @@ class LiberoEnv(gym.Env):
def close(self): def close(self):
if self._env is not None: if self._env is not None:
try: self._env.close()
self._env.close()
finally:
# LIBERO deletes its inner env on close, so this wrapper must
# be recreated before the next reset.
self._env = None
def _make_env_fns( def _make_env_fns(
+1 -3
View File
@@ -384,9 +384,7 @@ class RoboTwinEnv(gym.Env):
self._env: Any | None = None # deferred — created on first reset() inside worker self._env: Any | None = None # deferred — created on first reset() inside worker
self._step_count: int = 0 self._step_count: int = 0
self._black_frame: np.ndarray = np.zeros( self._black_frame = np.zeros((self.observation_height, self.observation_width, 3), dtype=np.uint8)
(self.observation_height, self.observation_width, 3), dtype=np.uint8
)
image_spaces = { image_spaces = {
cam: spaces.Box( cam: spaces.Box(
+1 -1
View File
@@ -373,7 +373,7 @@ class VLABenchEnv(gym.Env):
if action.shape[0] != 7: if action.shape[0] != 7:
# Unknown layout — fall back to zero-pad so the sim doesn't crash. # Unknown layout — fall back to zero-pad so the sim doesn't crash.
padded: np.ndarray = np.zeros(ctrl_dim, dtype=np.float64) padded = np.zeros(ctrl_dim, dtype=np.float64)
padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim] padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
return padded return padded
-18
View File
@@ -122,9 +122,6 @@ MODEL_ENCODING_TABLE = {
"xm430-w350": X_SERIES_ENCODINGS_TABLE, "xm430-w350": X_SERIES_ENCODINGS_TABLE,
"xm540-w270": X_SERIES_ENCODINGS_TABLE, "xm540-w270": X_SERIES_ENCODINGS_TABLE,
"xc430-w150": X_SERIES_ENCODINGS_TABLE, "xc430-w150": X_SERIES_ENCODINGS_TABLE,
"xh540-w150": X_SERIES_ENCODINGS_TABLE,
"xc330-t288": X_SERIES_ENCODINGS_TABLE,
"xc330-t181": X_SERIES_ENCODINGS_TABLE,
} }
# {model: model_resolution} # {model: model_resolution}
@@ -137,9 +134,6 @@ MODEL_RESOLUTION = {
"xm430-w350": 4096, "xm430-w350": 4096,
"xm540-w270": 4096, "xm540-w270": 4096,
"xc430-w150": 4096, "xc430-w150": 4096,
"xh540-w150": 4096,
"xc330-t288": 4096,
"xc330-t181": 4096,
} }
# {model: model_number} # {model: model_number}
@@ -151,9 +145,6 @@ MODEL_NUMBER_TABLE = {
"xm430-w350": 1020, "xm430-w350": 1020,
"xm540-w270": 1120, "xm540-w270": 1120,
"xc430-w150": 1070, "xc430-w150": 1070,
"xh540-w150": 1110,
"xc330-t288": 1220,
"xc330-t181": 1210,
} }
# {model: available_operating_modes} # {model: available_operating_modes}
@@ -165,9 +156,6 @@ MODEL_OPERATING_MODES = {
"xm430-w350": [0, 1, 3, 4, 5, 16], "xm430-w350": [0, 1, 3, 4, 5, 16],
"xm540-w270": [0, 1, 3, 4, 5, 16], "xm540-w270": [0, 1, 3, 4, 5, 16],
"xc430-w150": [1, 3, 4, 16], "xc430-w150": [1, 3, 4, 16],
"xh540-w150": [0, 1, 3, 4, 5, 16],
"xc330-t288": [0, 1, 3, 4, 5, 16],
"xc330-t181": [0, 1, 3, 4, 5, 16],
} }
MODEL_CONTROL_TABLE = { MODEL_CONTROL_TABLE = {
@@ -178,9 +166,6 @@ MODEL_CONTROL_TABLE = {
"xm430-w350": X_SERIES_CONTROL_TABLE, "xm430-w350": X_SERIES_CONTROL_TABLE,
"xm540-w270": X_SERIES_CONTROL_TABLE, "xm540-w270": X_SERIES_CONTROL_TABLE,
"xc430-w150": X_SERIES_CONTROL_TABLE, "xc430-w150": X_SERIES_CONTROL_TABLE,
"xh540-w150": X_SERIES_CONTROL_TABLE,
"xc330-t288": X_SERIES_CONTROL_TABLE,
"xc330-t181": X_SERIES_CONTROL_TABLE,
} }
MODEL_BAUDRATE_TABLE = { MODEL_BAUDRATE_TABLE = {
@@ -191,9 +176,6 @@ MODEL_BAUDRATE_TABLE = {
"xm430-w350": X_SERIES_BAUDRATE_TABLE, "xm430-w350": X_SERIES_BAUDRATE_TABLE,
"xm540-w270": X_SERIES_BAUDRATE_TABLE, "xm540-w270": X_SERIES_BAUDRATE_TABLE,
"xc430-w150": X_SERIES_BAUDRATE_TABLE, "xc430-w150": X_SERIES_BAUDRATE_TABLE,
"xh540-w150": X_SERIES_BAUDRATE_TABLE,
"xc330-t288": X_SERIES_BAUDRATE_TABLE,
"xc330-t181": X_SERIES_BAUDRATE_TABLE,
} }
AVAILABLE_BAUDRATES = [ AVAILABLE_BAUDRATES = [
+5 -35
View File
@@ -302,33 +302,6 @@ def _pad_evo1_stats(
return padded_stats return padded_stats
def _refresh_evo1_normalization_steps(
config: Evo1Config,
preprocessor: PolicyProcessorPipeline,
postprocessor: PolicyProcessorPipeline,
) -> None:
"""Re-pad checkpoint-loaded (un)normalizer stats/features to EVO1's fixed widths.
Loading a checkpoint injects the raw dataset stats (unpadded to max_state_dim/max_action_dim)
into the (un)normalizer via the generic override path in make_pre_post_processors. Those stats
and their declared features must be re-padded/reshaped to EVO1's fixed widths, otherwise
normalization fails against the padded state/action tensors (e.g. state padded to 24 vs. 8-dim
LIBERO stats). Padding is a no-op when stats are already at the target width.
"""
normalization_features = _evo1_normalization_features(config)
action_features = _evo1_action_features(config)
for step in preprocessor.steps:
if isinstance(step, NormalizerProcessorStep):
step.features = normalization_features
step.stats = _pad_evo1_stats(config, step.stats)
step.to(device=step.device, dtype=step.dtype)
for step in postprocessor.steps:
if isinstance(step, UnnormalizerProcessorStep):
step.features = action_features
step.stats = _pad_evo1_stats(config, step.stats)
step.to(device=step.device, dtype=step.dtype)
def reconcile_evo1_processors( def reconcile_evo1_processors(
config: Evo1Config, config: Evo1Config,
preprocessor: PolicyProcessorPipeline, preprocessor: PolicyProcessorPipeline,
@@ -336,19 +309,16 @@ def reconcile_evo1_processors(
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]: ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
"""Reconcile checkpoint-loaded pipelines with the current EVO1 config. """Reconcile checkpoint-loaded pipelines with the current EVO1 config.
Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
(converters are plain functions and are never serialized), eval-time CLI overrides of the (converters are plain functions and are never serialized), and eval-time CLI overrides of the
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This
(un)normalizer stats/features when the generic override path injects raw, unpadded dataset restores the converter and rebuilds the action step from the current config so those overrides
stats. This restores the converter, re-pads the normalization stats to EVO1's fixed widths, and take effect.
rebuilds the action step from the current config so those overrides take effect.
""" """
# Pipelines reloaded from a checkpoint come back with the default batch converter, which drops # Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
# non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1. # non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
preprocessor.to_transition = evo1_batch_to_transition preprocessor.to_transition = evo1_batch_to_transition
_refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
action_step = Evo1ActionProcessorStep( action_step = Evo1ActionProcessorStep(
action_dim=_evo1_action_dim(config), action_dim=_evo1_action_dim(config),
binarize_gripper=config.binarize_gripper, binarize_gripper=config.binarize_gripper,
+6 -18
View File
@@ -44,19 +44,12 @@ from lerobot.utils.constants import (
POLICY_PREPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME,
) )
from lerobot.utils.feature_utils import dataset_to_policy_features from lerobot.utils.feature_utils import dataset_to_policy_features
from lerobot.utils.import_utils import _peft_available, require_package
from .evo1.configuration_evo1 import Evo1Config from .evo1.configuration_evo1 import Evo1Config
from .groot.configuration_groot import GrootConfig from .groot.configuration_groot import GrootConfig
from .pretrained import PreTrainedPolicy from .pretrained import PreTrainedPolicy
from .utils import validate_visual_features_consistency from .utils import validate_visual_features_consistency
if TYPE_CHECKING or _peft_available:
from peft import PeftConfig, PeftModel
else:
PeftConfig = None
PeftModel = None
def _reconnect_relative_absolute_steps( def _reconnect_relative_absolute_steps(
preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
@@ -178,6 +171,9 @@ def make_pre_post_processors(
ValueError: If no processor factory exists for the given policy configuration type. ValueError: If no processor factory exists for the given policy configuration type.
""" """
if pretrained_path: if pretrained_path:
revision_resolver = getattr(policy_cfg, "get_hub_revision", None)
if callable(revision_resolver):
pretrained_revision = revision_resolver(pretrained_path, pretrained_revision)
if isinstance(policy_cfg, GrootConfig): if isinstance(policy_cfg, GrootConfig):
from .groot.processor_groot import make_groot_pre_post_processors_from_pretrained from .groot.processor_groot import make_groot_pre_post_processors_from_pretrained
@@ -341,15 +337,12 @@ def make_policy(
# Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo # Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
# of the adapter and the adapter's config contains the path to the base policy. So we need the # of the adapter and the adapter's config contains the path to the base policy. So we need the
# adapter config first, then load the correct policy and then apply PEFT. # adapter config first, then load the correct policy and then apply PEFT.
require_package("peft", extra="peft") from peft import PeftConfig, PeftModel
logging.info("Loading policy's PEFT adapter.") logging.info("Loading policy's PEFT adapter.")
peft_pretrained_path = str(cfg.pretrained_path) peft_pretrained_path = str(cfg.pretrained_path)
peft_config = PeftConfig.from_pretrained( peft_config = PeftConfig.from_pretrained(peft_pretrained_path)
peft_pretrained_path,
revision=cfg.pretrained_revision,
)
kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path
if not kwargs["pretrained_name_or_path"]: if not kwargs["pretrained_name_or_path"]:
@@ -360,14 +353,9 @@ def make_policy(
"the adapter was trained." "the adapter was trained."
) )
kwargs["revision"] = peft_config.revision
policy = policy_cls.from_pretrained(**kwargs) policy = policy_cls.from_pretrained(**kwargs)
policy = PeftModel.from_pretrained( policy = PeftModel.from_pretrained(
policy, policy, peft_pretrained_path, config=peft_config, is_trainable=True
peft_pretrained_path,
config=peft_config,
revision=cfg.pretrained_revision,
is_trainable=True,
) )
else: else:
@@ -37,19 +37,13 @@ def is_image_feature(key: str) -> bool:
@dataclass @dataclass
class ConcurrencyConfig: class ConcurrencyConfig:
"""Configuration for the concurrency of the actor and learner. """Configuration for the concurrency of the actor and learner.
Possible values are: Possible values are:
- "threads": Use threads for the actor and learner. - "threads": Use threads for the actor and learner.
- "processes": Use processes for the actor and learner. - "processes": Use processes for the actor and learner.
``multiprocessing_context`` selects the process-wide start method when
processes are used. Set it to ``None`` to preserve Python's default or a
method already selected by the embedding application.
""" """
actor: str = "threads" actor: str = "threads"
learner: str = "threads" learner: str = "threads"
multiprocessing_context: str | None = "spawn"
@dataclass @dataclass
+8 -1
View File
@@ -37,6 +37,7 @@ from torch import Tensor
from lerobot.configs import FeatureType, PolicyFeature from lerobot.configs import FeatureType, PolicyFeature
from lerobot.utils.constants import ACTION, OBS_IMAGES from lerobot.utils.constants import ACTION, OBS_IMAGES
from lerobot.utils.hub import extract_commit_hash
from lerobot.utils.import_utils import _transformers_available, require_package from lerobot.utils.import_utils import _transformers_available, require_package
from ..pretrained import PreTrainedPolicy from ..pretrained import PreTrainedPolicy
@@ -195,6 +196,8 @@ class GrootPolicy(PreTrainedPolicy):
) )
model_id = str(pretrained_name_or_path) model_id = str(pretrained_name_or_path)
if config is not None:
revision = config.get_hub_revision(model_id, revision)
is_finetuned_checkpoint = False is_finetuned_checkpoint = False
# Check if this is a fine-tuned LeRobot checkpoint (has model.safetensors) # Check if this is a fine-tuned LeRobot checkpoint (has model.safetensors)
@@ -204,7 +207,7 @@ class GrootPolicy(PreTrainedPolicy):
else: else:
# Try to download the safetensors file to check if it exists # Try to download the safetensors file to check if it exists
try: try:
hf_hub_download( resolved_model_file = hf_hub_download(
repo_id=model_id, repo_id=model_id,
filename=SAFETENSORS_SINGLE_FILE, filename=SAFETENSORS_SINGLE_FILE,
revision=revision, revision=revision,
@@ -214,6 +217,10 @@ class GrootPolicy(PreTrainedPolicy):
token=token, token=token,
local_files_only=local_files_only, local_files_only=local_files_only,
) )
resolved_commit_hash = extract_commit_hash(resolved_model_file, revision)
revision = resolved_commit_hash or revision
if config is not None and config._commit_hash is None:
config._set_hub_commit_hash(resolved_commit_hash, model_id)
is_finetuned_checkpoint = True is_finetuned_checkpoint = True
except HfHubHTTPError: except HfHubHTTPError:
is_finetuned_checkpoint = False is_finetuned_checkpoint = False
@@ -43,22 +43,11 @@ from torch.distributions import Beta
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.utils.constants import ACTION from lerobot.utils.constants import ACTION
from lerobot.utils.import_utils import ( from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
_peft_available,
_scipy_available,
_transformers_available,
require_package,
)
from ..rtc.modeling_rtc import RTCProcessor from ..rtc.modeling_rtc import RTCProcessor
from .configuration_molmoact2 import MolmoAct2Config from .configuration_molmoact2 import MolmoAct2Config
if TYPE_CHECKING or _peft_available:
from peft import LoraConfig, get_peft_model
else:
LoraConfig = None
get_peft_model = None
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -1742,11 +1731,13 @@ class MolmoAct2Policy(PreTrainedPolicy):
def _build_inner_lora_config(self): def _build_inner_lora_config(self):
require_package("peft", extra="molmoact2") require_package("peft", extra="molmoact2")
from peft import LoraConfig
return LoraConfig(**self._get_inner_peft_targets()) return LoraConfig(**self._get_inner_peft_targets())
def _apply_lora_adapters(self) -> None: def _apply_lora_adapters(self) -> None:
require_package("peft", extra="molmoact2") require_package("peft", extra="molmoact2")
from peft import get_peft_model
peft_config = self._build_inner_lora_config() peft_config = self._build_inner_lora_config()
self._validate_peft_config(peft_config) self._validate_peft_config(peft_config)
+8 -1
View File
@@ -53,6 +53,7 @@ from lerobot.utils.constants import (
OBS_LANGUAGE_TOKENS, OBS_LANGUAGE_TOKENS,
OBS_STATE, OBS_STATE,
) )
from lerobot.utils.hub import extract_commit_hash
from ..common.flow_matching import euler_integrate, sample_noise, sample_time_beta from ..common.flow_matching import euler_integrate, sample_noise, sample_time_beta
from ..common.vla_utils import ( from ..common.vla_utils import (
@@ -814,6 +815,8 @@ class PI0Policy(PreTrainedPolicy):
**kwargs, **kwargs,
) )
revision = config.get_hub_revision(pretrained_name_or_path, revision)
# Initialize model without loading weights # Initialize model without loading weights
# Check if dataset_stats were provided in kwargs # Check if dataset_stats were provided in kwargs
model = cls(config, **kwargs) model = cls(config, **kwargs)
@@ -832,9 +835,13 @@ class PI0Policy(PreTrainedPolicy):
resume_download=kwargs.get("resume_download"), resume_download=kwargs.get("resume_download"),
proxies=kwargs.get("proxies"), proxies=kwargs.get("proxies"),
token=kwargs.get("token"), token=kwargs.get("token"),
revision=kwargs.get("revision"), revision=revision,
local_files_only=kwargs.get("local_files_only", False), local_files_only=kwargs.get("local_files_only", False),
) )
if config._commit_hash is None:
config._set_hub_commit_hash(
extract_commit_hash(resolved_file, revision), str(pretrained_name_or_path)
)
from safetensors.torch import load_file from safetensors.torch import load_file
original_state_dict = load_file(resolved_file) original_state_dict = load_file(resolved_file)
+8 -1
View File
@@ -50,6 +50,7 @@ from lerobot.utils.constants import (
OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_ATTENTION_MASK,
OBS_LANGUAGE_TOKENS, OBS_LANGUAGE_TOKENS,
) )
from lerobot.utils.hub import extract_commit_hash
from ..common.flow_matching import euler_integrate, sample_noise, sample_time_beta from ..common.flow_matching import euler_integrate, sample_noise, sample_time_beta
from ..common.vla_utils import ( from ..common.vla_utils import (
@@ -779,6 +780,8 @@ class PI05Policy(PreTrainedPolicy):
**kwargs, **kwargs,
) )
revision = config.get_hub_revision(pretrained_name_or_path, revision)
# Initialize model without loading weights # Initialize model without loading weights
# Check if dataset_stats were provided in kwargs # Check if dataset_stats were provided in kwargs
model = cls(config, **kwargs) model = cls(config, **kwargs)
@@ -797,9 +800,13 @@ class PI05Policy(PreTrainedPolicy):
resume_download=kwargs.get("resume_download"), resume_download=kwargs.get("resume_download"),
proxies=kwargs.get("proxies"), proxies=kwargs.get("proxies"),
token=kwargs.get("token"), token=kwargs.get("token"),
revision=kwargs.get("revision"), revision=revision,
local_files_only=kwargs.get("local_files_only", False), local_files_only=kwargs.get("local_files_only", False),
) )
if config._commit_hash is None:
config._set_hub_commit_hash(
extract_commit_hash(resolved_file, revision), str(pretrained_name_or_path)
)
from safetensors.torch import load_file from safetensors.torch import load_file
original_state_dict = load_file(resolved_file) original_state_dict = load_file(resolved_file)
@@ -55,6 +55,7 @@ from lerobot.utils.constants import (
OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_ATTENTION_MASK,
OBS_LANGUAGE_TOKENS, OBS_LANGUAGE_TOKENS,
) )
from lerobot.utils.hub import extract_commit_hash
from ..common.vla_utils import pad_vector, prepare_attention_masks_4d, resize_with_pad_torch from ..common.vla_utils import pad_vector, prepare_attention_masks_4d, resize_with_pad_torch
from ..pretrained import PreTrainedPolicy, T from ..pretrained import PreTrainedPolicy, T
@@ -808,6 +809,8 @@ class PI0FastPolicy(PreTrainedPolicy):
**kwargs, **kwargs,
) )
revision = config.get_hub_revision(pretrained_name_or_path, revision)
# Initialize model without loading weights # Initialize model without loading weights
# Check if dataset_stats were provided in kwargs # Check if dataset_stats were provided in kwargs
model = cls(config, **kwargs) model = cls(config, **kwargs)
@@ -826,9 +829,13 @@ class PI0FastPolicy(PreTrainedPolicy):
resume_download=kwargs.get("resume_download"), resume_download=kwargs.get("resume_download"),
proxies=kwargs.get("proxies"), proxies=kwargs.get("proxies"),
token=kwargs.get("token"), token=kwargs.get("token"),
revision=kwargs.get("revision"), revision=revision,
local_files_only=kwargs.get("local_files_only", False), local_files_only=kwargs.get("local_files_only", False),
) )
if config._commit_hash is None:
config._set_hub_commit_hash(
extract_commit_hash(resolved_file, revision), str(pretrained_name_or_path)
)
from safetensors.torch import load_file from safetensors.torch import load_file
original_state_dict = load_file(resolved_file) original_state_dict = load_file(resolved_file)
+9 -14
View File
@@ -33,23 +33,15 @@ from lerobot.__version__ import __version__
from lerobot.configs import PreTrainedConfig from lerobot.configs import PreTrainedConfig
from lerobot.configs.train import TrainPipelineConfig from lerobot.configs.train import TrainPipelineConfig
from lerobot.utils.device_utils import resolve_safetensors_device from lerobot.utils.device_utils import resolve_safetensors_device
from lerobot.utils.hub import HubMixin from lerobot.utils.hub import HubMixin, extract_commit_hash
from lerobot.utils.import_utils import _peft_available, require_package
from .utils import log_model_loading_keys from .utils import log_model_loading_keys
if TYPE_CHECKING or _peft_available: T = TypeVar("T", bound="PreTrainedPolicy")
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType, get_peft_model
else:
PEFT_TYPE_TO_CONFIG_MAPPING = None
PeftType = None
get_peft_model = None
if TYPE_CHECKING: if TYPE_CHECKING:
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
T = TypeVar("T", bound="PreTrainedPolicy")
def _build_card_context( def _build_card_context(
cfg: TrainPipelineConfig | None, cfg: TrainPipelineConfig | None,
@@ -198,6 +190,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
**kwargs, **kwargs,
) )
model_id = str(pretrained_name_or_path) model_id = str(pretrained_name_or_path)
revision = config.get_hub_revision(model_id, revision)
instance = cls(config, **kwargs) instance = cls(config, **kwargs)
if os.path.isdir(model_id): if os.path.isdir(model_id):
print("Loading weights from local directory") print("Loading weights from local directory")
@@ -216,6 +209,8 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
token=token, token=token,
local_files_only=local_files_only, local_files_only=local_files_only,
) )
if config._commit_hash is None:
config._set_hub_commit_hash(extract_commit_hash(model_file, revision), model_id)
policy = cls._load_as_safetensor(instance, model_file, config.device, strict) policy = cls._load_as_safetensor(instance, model_file, config.device, strict)
except HfHubHTTPError as e: except HfHubHTTPError as e:
raise FileNotFoundError( raise FileNotFoundError(
@@ -392,7 +387,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.) peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.)
These are merged with policy defaults to build the final config. These are merged with policy defaults to build the final config.
""" """
require_package("peft", extra="peft") from peft import get_peft_model
# If user provided a complete config, use it directly (with overrides) # If user provided a complete config, use it directly (with overrides)
if peft_config is not None: if peft_config is not None:
@@ -463,7 +458,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
Returns: Returns:
Preprocessed dict with renamed keys and init_type mapped to method-specific key. Preprocessed dict with renamed keys and init_type mapped to method-specific key.
""" """
require_package("peft", extra="peft") from peft import PeftType
cli_overrides = cli_overrides.copy() cli_overrides = cli_overrides.copy()
@@ -488,7 +483,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
def _build_peft_config(self, cli_overrides: dict): def _build_peft_config(self, cli_overrides: dict):
"""Build a PEFT config from policy defaults and CLI overrides.""" """Build a PEFT config from policy defaults and CLI overrides."""
require_package("peft", extra="peft") from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
# Determine PEFT method type (default to LORA) # Determine PEFT method type (default to LORA)
method_type_str = cli_overrides.get("method_type") or "lora" method_type_str = cli_overrides.get("method_type") or "lora"
@@ -515,7 +510,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict): def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
"""Apply CLI overrides to an existing PEFT config.""" """Apply CLI overrides to an existing PEFT config."""
require_package("peft", extra="peft") from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
# Get method type from existing config or CLI override # Get method type from existing config or CLI override
method_type_str = cli_overrides.get("method_type") method_type_str = cli_overrides.get("method_type")
@@ -31,6 +31,7 @@ from torch import Tensor, nn
from lerobot.configs import PreTrainedConfig from lerobot.configs import PreTrainedConfig
from lerobot.utils.constants import ACTION, OBS_LANGUAGE_TOKENS, OBS_STATE from lerobot.utils.constants import ACTION, OBS_LANGUAGE_TOKENS, OBS_STATE
from lerobot.utils.hub import extract_commit_hash
from lerobot.utils.import_utils import _transformers_available, require_package from lerobot.utils.import_utils import _transformers_available, require_package
from ..common.vla_utils import pad_vector, resize_with_pad from ..common.vla_utils import pad_vector, resize_with_pad
@@ -459,6 +460,7 @@ class XVLAPolicy(PreTrainedPolicy):
) )
model_id = str(pretrained_name_or_path) model_id = str(pretrained_name_or_path)
revision = config.get_hub_revision(model_id, revision)
instance = cls(config, **kwargs) instance = cls(config, **kwargs)
# step 2: locate model.safetensors # step 2: locate model.safetensors
if os.path.isdir(model_id): if os.path.isdir(model_id):
@@ -480,6 +482,8 @@ class XVLAPolicy(PreTrainedPolicy):
token=token, token=token,
local_files_only=local_files_only, local_files_only=local_files_only,
) )
if config._commit_hash is None:
config._set_hub_commit_hash(extract_commit_hash(model_file, revision), model_id)
except HfHubHTTPError as e: except HfHubHTTPError as e:
raise FileNotFoundError(f"model.safetensors not found on the Hub at {model_id}") from e raise FileNotFoundError(f"model.safetensors not found on the Hub at {model_id}") from e
+5 -1
View File
@@ -47,7 +47,7 @@ from safetensors.torch import load_file, save_file
from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.types import EnvAction, EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey from lerobot.types import EnvAction, EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
from lerobot.utils.constants import HF_LEROBOT_HOME from lerobot.utils.constants import HF_LEROBOT_HOME
from lerobot.utils.hub import HubMixin from lerobot.utils.hub import HubMixin, extract_commit_hash
from .converters import batch_to_transition, create_transition, transition_to_batch from .converters import batch_to_transition, create_transition, transition_to_batch
@@ -727,6 +727,10 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
# 1. Load configuration using simplified 3-way logic # 1. Load configuration using simplified 3-way logic
loaded_config, base_path = cls._load_config(model_id, config_filename, hub_download_kwargs) loaded_config, base_path = cls._load_config(model_id, config_filename, hub_download_kwargs)
if not is_local_source:
commit_hash = extract_commit_hash(base_path, revision)
if commit_hash is not None:
hub_download_kwargs["revision"] = commit_hash
# 2. Validate configuration and handle migration # 2. Validate configuration and handle migration
cls._validate_loaded_config(model_id, loaded_config, config_filename) cls._validate_loaded_config(model_id, loaded_config, config_filename)
+4 -2
View File
@@ -91,7 +91,7 @@ from lerobot.robots import so_follower # noqa: F401
from lerobot.teleoperators import gamepad, so_leader # noqa: F401 from lerobot.teleoperators import gamepad, so_leader # noqa: F401
from lerobot.teleoperators.utils import TeleopEvents from lerobot.teleoperators.utils import TeleopEvents
from lerobot.utils.device_utils import get_safe_torch_device from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method from lerobot.utils.process import ProcessSignalHandler
from lerobot.utils.random_utils import set_seed from lerobot.utils.random_utils import set_seed
from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.transition import ( from lerobot.utils.transition import (
@@ -124,7 +124,9 @@ def actor_cli(cfg: TrainRLServerPipelineConfig):
cfg.validate() cfg.validate()
display_pid = False display_pid = False
if not use_threads(cfg): if not use_threads(cfg):
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context) import torch.multiprocessing as mp
mp.set_start_method("spawn")
display_pid = True display_pid = True
# Create logs directory to ensure it exists # Create logs directory to ensure it exists
+2 -2
View File
@@ -18,7 +18,7 @@ import functools
import threading import threading
from collections.abc import Callable, Sequence from collections.abc import Callable, Sequence
from contextlib import suppress from contextlib import suppress
from typing import NotRequired, TypedDict from typing import TypedDict
import torch import torch
import torch.nn.functional as F # noqa: N812 import torch.nn.functional as F # noqa: N812
@@ -36,7 +36,7 @@ class BatchTransition(TypedDict):
next_state: dict[str, torch.Tensor] next_state: dict[str, torch.Tensor]
done: torch.Tensor done: torch.Tensor
truncated: torch.Tensor truncated: torch.Tensor
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None] complementary_info: dict[str, torch.Tensor | float | int] | None = None
def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor: def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
+4 -2
View File
@@ -102,7 +102,7 @@ from lerobot.utils.constants import (
) )
from lerobot.utils.device_utils import get_safe_torch_device from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.io_utils import load_json, write_json from lerobot.utils.io_utils import load_json, write_json
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method from lerobot.utils.process import ProcessSignalHandler
from lerobot.utils.random_utils import set_seed from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import ( from lerobot.utils.utils import (
format_big_number, format_big_number,
@@ -123,7 +123,9 @@ def train_cli(cfg: TrainRLServerPipelineConfig):
# Fail fast with a friendly error if the optional ``hilserl`` extra is missing. # Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
require_package("grpcio", extra="hilserl", import_name="grpc") require_package("grpcio", extra="hilserl", import_name="grpc")
if not use_threads(cfg): if not use_threads(cfg):
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context) import torch.multiprocessing as mp
mp.set_start_method("spawn")
# Use the job_name from the config # Use the job_name from the config
train( train(
@@ -46,12 +46,6 @@ class SOFollowerConfig:
position_i_coefficient: int = 0 position_i_coefficient: int = 0
position_d_coefficient: int = 32 position_d_coefficient: int = 32
# Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
# return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
# at once, which otherwise aborts the control loop. Retries are immediate (no sleep) and only happen on
# failure, so the steady-state read cost is unchanged.
num_read_retries: int = 2
@RobotConfig.register_subclass("so101_follower") @RobotConfig.register_subclass("so101_follower")
@RobotConfig.register_subclass("so100_follower") @RobotConfig.register_subclass("so100_follower")
@@ -510,10 +510,10 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
# We only use the ee pose in the dataset, so we don't need the joint positions # We only use the ee pose in the dataset, so we don't need the joint positions
for n in self.motor_names: for n in self.motor_names:
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None) features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
# Store end-effector features as actions in the dataset schema # We specify the dataset features of this step that we want to be stored in the dataset
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]: for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature( features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
type=FeatureType.ACTION, shape=(1,) type=FeatureType.STATE, shape=(1,)
) )
return features return features
@@ -180,7 +180,7 @@ class SOFollower(Robot):
def get_observation(self) -> RobotObservation: def get_observation(self) -> RobotObservation:
# Read arm position # Read arm position
start = time.perf_counter() start = time.perf_counter()
obs_dict = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries) obs_dict = self.bus.sync_read("Present_Position")
obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()} obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
dt_ms = (time.perf_counter() - start) * 1e3 dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read state: {dt_ms:.1f}ms") logger.debug(f"{self} read state: {dt_ms:.1f}ms")
@@ -221,7 +221,7 @@ class SOFollower(Robot):
# Cap goal position when too far away from present position. # Cap goal position when too far away from present position.
# /!\ Slower fps expected due to reading from the follower. # /!\ Slower fps expected due to reading from the follower.
if self.config.max_relative_target is not None: if self.config.max_relative_target is not None:
present_pos = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries) present_pos = self.bus.sync_read("Present_Position")
goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()} goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()}
goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target) goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)
@@ -68,10 +68,6 @@ class UnitreeG1Config(RobotConfig):
# Compensates for gravity on the unitree's arms using the arm ik solver # Compensates for gravity on the unitree's arms using the arm ik solver
gravity_compensation: bool = False gravity_compensation: bool = False
# Locomotion controller class name, e.g. "GrootLocomotionController", # Lower-body controller class name, e.g. "GrootLocomotionController" or
# "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it. # "HolosomaLocomotionController". None disables it.
# Selecting "SonicWholeBodyController" implicitly switches the robot to the 64-D
# latent-token action/observation interface (``motion_token.{i}.pos`` action and a
# ``motion_token_state.{i}.pos`` state echo) so ``lerobot-rollout`` can drive a
# policy trained on SONIC motion tokens (e.g. nepyope/sonic_walk).
controller: str | None = None controller: str | None = None
@@ -1,27 +0,0 @@
#!/usr/bin/env python
# Copyright 2025 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.
"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
from .gr00t_locomotion import GrootLocomotionController
from .holosoma_locomotion import HolosomaLocomotionController
from .sonic_whole_body import SonicWholeBodyController
__all__ = [
"GrootLocomotionController",
"HolosomaLocomotionController",
"SonicWholeBodyController",
]
@@ -1,378 +0,0 @@
#!/usr/bin/env python
# Copyright 2025 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.
"""SONIC decoder whole-body controller for the Unitree G1 (token-only).
Pure-Python/ONNX re-implementation of the *decode* half of NVIDIA's SONIC deploy stack.
The encoder is intentionally absent: a token-output VLA (e.g. ``nepyope/sonic_walk``)
supplies the 64-D latent ``motion_token`` directly each tick, and the SONIC **decoder**
maps ``token + recent proprioception history`` to a residual action that is scaled and
added onto the standing pose (``default_angles``) to produce 50 Hz joint-position targets
for the robot's PD controller.
Index spaces: joints exist in two orderings **IsaacLab** (policy/training order) and
**MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` (in g1_utils)
convert between them. Quaternions are scalar-first ``(w, x, y, z)``.
"""
from __future__ import annotations
import json
import logging
import numpy as np
import onnx
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from ..g1_utils import (
ISAACLAB_TO_MUJOCO,
MUJOCO_TO_ISAACLAB,
G1_29_JointIndex,
get_gravity_orientation,
)
from ..unitree_g1 import lowstate_to_obs
logger = logging.getLogger(__name__)
# ── Constants (hardware-validated; see the NVIDIA SONIC deploy reference) ──────
CONTROL_DT = 0.02 # 50 Hz control period (s)
TOKEN_DIM = 64 # decoder latent size
# SONIC decoder checkpoint: NVIDIA's decoder ONNX re-packaged with its deploy constants
# (kp/kd PD gains, the standing pose default_angles, and the residual action_scale) embedded
# in the ONNX metadata; see upload_sonic_decoder.py for provisioning. The runtime loads the
# model *and* all of these straight from the checkpoint (the Holosoma convention), so no
# motor-physics math happens at deploy time.
DEFAULT_SONIC_REPO_ID = "lerobot/sonic_decoder"
DECODER_FILENAME = "model_decoder.onnx"
DECODER_INPUT_DIM = 994 # token(64) + 10-frame proprio history + gravity
def load_sonic_decoder(repo_id: str = DEFAULT_SONIC_REPO_ID):
"""Load the SONIC decoder ONNX and its baked-in deploy constants from the checkpoint.
Returns ``(decoder_session, kp, kd, default_angles, action_scale, neutral_token)``. The
gains/pose/scale are (29,) float32 in IsaacLab joint order and ``neutral_token`` is the
(64,) float32 idle latent -- all read from the ONNX ``metadata_props`` rather than
recomputed/hardcoded at deploy time (mirrors ``holosoma_locomotion.load_policy``).
"""
decoder_path = hf_hub_download(repo_id=repo_id, filename=DECODER_FILENAME)
so = ort.SessionOptions()
so.log_severity_level = 3 # quiet ORT logs
session = ort.InferenceSession(decoder_path, sess_options=so)
dec_dim = int(session.get_inputs()[0].shape[1])
if dec_dim != DECODER_INPUT_DIM:
raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected {DECODER_INPUT_DIM})")
meta = {p.key: p.value for p in onnx.load(decoder_path, load_external_data=False).metadata_props}
required = ("kp", "kd", "default_angles", "action_scale", "neutral_token")
missing = [k for k in required if k not in meta]
if missing:
raise ValueError(
f"SONIC decoder ONNX at {repo_id} is missing metadata {missing}; "
"re-run upload_sonic_decoder.py to (re)provision the checkpoint."
)
arr = {k: np.array(json.loads(meta[k]), dtype=np.float32) for k in required}
logger.info("Loaded SONIC deploy constants from %s (%d joints)", repo_id, len(arr["kp"]))
return session, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"]
def _to_mujoco(a):
"""Apply the ``MUJOCO_TO_ISAACLAB`` gather to a 29-vector (deploy-order reorder).
NOTE: this returns ``a[MUJOCO_TO_ISAACLAB]``. The ``_mj`` suffixes and the exact
permutation direction are a fixed convention validated against the deployed SONIC ONNX
policy (the decoder consumes vectors in this order). Do not "correct" the table or
rename toward the opposite direction without re-validating on hardware.
"""
return a[MUJOCO_TO_ISAACLAB]
# Action-feature prefix for the latent-token interface (see _extract_token_from_action).
TOKEN_ACTION_PREFIX = "motion_token" # nosec B105 - feature-key prefix, not a secret
# Proprio-state prefix for the token interface: the robot echoes the last commanded token
# here so ``lerobot-rollout`` aggregates it into a 64-D ``observation.state``.
TOKEN_STATE_PREFIX = "motion_token_state" # nosec B105 - feature-key prefix, not a secret
def token_action_key(i: int) -> str:
"""Action-dict key for the i-th component of the 64-D SONIC latent token.
The ``.pos`` suffix is required so the value flows through ``lerobot-rollout``, which
only routes ``.pos`` scalar features onto the policy action vector.
"""
return f"{TOKEN_ACTION_PREFIX}.{i}.pos"
def token_state_key(i: int) -> str:
"""Observation key for the i-th component of the 64-D SONIC latent token state."""
return f"{TOKEN_STATE_PREFIX}.{i}.pos"
# Startup blend duration: over the first control ticks, linearly interpolate every joint
# from the robot's initial measured pose into the policy's commanded target, so control
# eases in without a snap on the first command.
INIT_RAMP_S = 3.0
def _extract_token_from_action(action: dict | None) -> np.ndarray | None:
"""Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None.
The token-only interface: the caller supplies the 64-D encoder latent directly (e.g. a
token-output VLA's action), which the decoder consumes with the encoder bypassed.
Requires the full dense token; a partial one is ignored (returns None).
"""
if not action:
return None
keys = [token_action_key(i) for i in range(TOKEN_DIM)]
if any(key not in action for key in keys):
return None
return np.fromiter((float(action[key]) for key in keys), dtype=np.float32, count=TOKEN_DIM)
class SonicDecoder:
"""Runs the SONIC decoder ONNX model and owns the proprioception history.
Each tick it appends the latest robot state to 10-frame history buffers, then maps the
supplied 64-D ``token`` + that history to a residual action added onto ``default_angles``.
The encoder is bypassed entirely (token supplied by the policy). ``default_angles`` and
``action_scale`` are (29,) float32 in IsaacLab order, loaded from the checkpoint.
"""
def __init__(self, decoder, default_angles, action_scale):
self.decoder = decoder
self.decoder_input = decoder.get_inputs()[0].name
self.default_angles = np.asarray(default_angles, np.float32)
self.action_scale = np.asarray(action_scale, np.float32)
self.default_angles_mj = _to_mujoco(self.default_angles)
self.token = np.zeros(TOKEN_DIM, np.float32)
self.last_action_mj = np.zeros(29, np.float32)
self.h_q_mj = [np.zeros(29, np.float32)] * 10
self.h_dq_mj = [np.zeros(29, np.float32)] * 10
self.h_ang = [np.zeros(3, np.float32)] * 10
self.h_act_mj = [np.zeros(29, np.float32)] * 10
self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
def reset(self):
"""Clear the token and 10-frame proprioception history.
``UnitreeG1.reset()`` relies on this so the first decoder outputs of a new episode
are not contaminated by the previous episode's state.
"""
self.token = np.zeros(TOKEN_DIM, np.float32)
self.last_action_mj = np.zeros(29, np.float32)
self.h_q_mj = [np.zeros(29, np.float32)] * 10
self.h_dq_mj = [np.zeros(29, np.float32)] * 10
self.h_ang = [np.zeros(3, np.float32)] * 10
self.h_act_mj = [np.zeros(29, np.float32)] * 10
self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
def update_history(self, q, dq, ang, quat):
"""Push the latest proprioception (pos/vel/gyro/orientation) into the 10-frame buffers."""
quat = quat / (np.linalg.norm(quat) + 1e-8)
q_mj = _to_mujoco(q)
dq_mj = _to_mujoco(dq)
self.h_q_mj = [q_mj - self.default_angles_mj] + self.h_q_mj[:-1]
self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
self.h_ang = [ang.copy()] + self.h_ang[:-1]
self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
self.h_quat = [quat.copy()] + self.h_quat[:-1]
def build_decoder_obs(self):
"""Assemble the 994-D decoder input: token + 10-frame proprioception history + gravity."""
obs = np.zeros(994, np.float32)
off = 0
obs[off : off + 64] = self.token
off += 64
for h, sz in [
(list(reversed(self.h_ang)), 3),
(list(reversed(self.h_q_mj)), 29),
(list(reversed(self.h_dq_mj)), 29),
(list(reversed(self.h_act_mj)), 29),
]:
for f in range(10):
obs[off : off + sz] = h[f]
off += sz
for q in reversed(self.h_quat):
obs[off : off + 3] = get_gravity_orientation(q)
off += 3
assert off == 994, f"Decoder obs mismatch: {off}"
return obs
def step(self, robot_obs, token, debug=False):
"""One control tick: read robot obs, decode the supplied token -> joint targets.
Args:
robot_obs: dict with ``<joint>.q``/``.dq`` and ``imu.*`` fields.
token: 64-D latent supplied by the policy (encoder bypassed).
debug: log action/delta norms.
Returns:
dict of ``<joint>.q`` target positions (rad) in IsaacLab joint order.
"""
self.token = np.asarray(token, np.float32)
jnames = [m.name for m in G1_29_JointIndex]
q = np.array(
[
robot_obs.get(f"{n}.q", self.default_angles[m.value])
for m, n in zip(G1_29_JointIndex, jnames, strict=False)
],
np.float32,
)
dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32)
quat = np.array(
[
robot_obs.get("imu.quat.w", 1),
robot_obs.get("imu.quat.x", 0),
robot_obs.get("imu.quat.y", 0),
robot_obs.get("imu.quat.z", 0),
],
np.float32,
)
ang = np.array([robot_obs.get(f"imu.gyro.{a}", 0) for a in "xyz"], np.float32)
self.update_history(q, dq, ang, quat)
action_mj = (
self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1, -1)})[0]
.squeeze()
.astype(np.float32)
)
self.last_action_mj = action_mj.copy()
target = self.default_angles + action_mj[ISAACLAB_TO_MUJOCO] * self.action_scale
if debug:
delta = target - q
logger.debug(
"token_norm=%.4f action_norm=%.4f delta_max=%.4f delta_rms=%.4f",
np.linalg.norm(self.token),
np.linalg.norm(action_mj),
np.max(np.abs(delta)),
np.sqrt(np.mean(delta**2)),
)
return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
class SonicRuntime:
"""Loads the SONIC decoder ONNX model and owns the decode controller.
Token-only deploy: the encoder is bypassed; each tick the decoder consumes a 64-D
latent token supplied directly by the policy.
"""
def __init__(self):
decoder_sess, self.kp, self.kd, default_angles, action_scale, neutral_token = load_sonic_decoder()
self.default_angles = default_angles
self.neutral_token = neutral_token
self.controller = SonicDecoder(decoder_sess, default_angles, action_scale)
@property
def pipeline(self):
return self.controller
def reset(self):
self.controller.reset()
def shutdown(self):
pass
class SonicWholeBodyController:
"""Full-body SONIC controller for UnitreeG1's background controller thread."""
control_dt = CONTROL_DT
full_body = True
def __init__(self):
logger.info("Loading SONIC whole-body controller...")
self._runtime = SonicRuntime()
self.kp = self._runtime.kp
self.kd = self._runtime.kd
self.controller = self._runtime.controller
self._default_angles = self._runtime.default_angles
self._neutral_token = self._runtime.neutral_token
# Startup blend: ease from the robot's initial pose into the first commanded policy
# targets over INIT_RAMP_S (captured on the first control tick).
self._init_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
self._init_step = 0
self._start_pose: dict[str, float] = {}
# Token-interface state. ``token_mode`` is set True by the robot whenever a SONIC
# whole-body controller is selected (token-driven deploy): the controller then holds a
# stable *neutral* token until the first real token arrives, and afterwards holds the
# *last* token received between ticks (the async controller runs ~50 Hz while a token
# VLA streams ~30 Hz). This lives here (not in the entry-point script) so it applies
# uniformly to run_g1_server, lerobot-rollout and the sim replays.
self.token_mode = False
self._last_token: np.ndarray | None = None
logger.info("SONIC ready (decoder, 64-D token command path)")
def _startup_blend(self, obs: dict, out: dict) -> dict:
"""Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds,
interpolate between the robot's pose captured on the first tick and the policy's
live commanded target, so the handoff has no snap.
``out`` is the policy's ``<joint>.q`` target dict for this tick; the blend ratio
climbs 0->1 over the ramp, after which the raw policy target passes through.
"""
if self._init_step >= self._init_ramp_steps or not out:
return out
if self._init_step == 0:
# Capture the robot's actual pose as the interpolation start point.
self._start_pose = {
f"{m.name}.q": float(obs.get(f"{m.name}.q", self._default_angles[m.value]))
for m in G1_29_JointIndex
}
self._init_step += 1
ratio = min(1.0, self._init_step / self._init_ramp_steps)
blended = {
k: self._start_pose.get(k, float(tgt)) * (1.0 - ratio) + float(tgt) * ratio
for k, tgt in out.items()
}
if self._init_step >= self._init_ramp_steps:
logger.info("SONIC startup blend complete -> full policy control")
return blended
def run_step(self, action: dict, lowstate) -> dict:
if lowstate is None:
return {}
obs = lowstate_to_obs(lowstate)
# Token-only interface (token-output VLA): a dense 64-D ``motion_token.{i}`` command
# is decoded directly, encoder bypassed.
token = _extract_token_from_action(action)
if token is not None:
self._last_token = token
elif self._last_token is None and self.token_mode:
# Token-driven deploy, but no token has arrived yet: hold the checkpoint's neutral
# token, which the decoder maps to a stable, natural standing pose.
self._last_token = self._neutral_token.copy()
if self._last_token is None:
# No token yet and not in token_mode: hold (keep last target).
return {}
# Either a fresh token this tick or the last one received (held between the ~30 Hz
# token stream and the ~50 Hz control loop).
return self._startup_blend(obs, self.controller.step(obs, self._last_token))
def reset(self):
self._runtime.reset()
self._init_step = 0 # re-run the startup blend after a reset
self._start_pose = {}
# Drop the held token so token_mode re-seeds the neutral token after a reset.
self._last_token = None
def shutdown(self):
self._runtime.shutdown()
+2 -44
View File
@@ -23,47 +23,6 @@ import numpy as np
NUM_MOTORS = 29 NUM_MOTORS = 29
# Joint-order permutations between the two 29-DoF layouts used across the G1 stack:
# IsaacLab (policy/training order) and MuJoCo (deploy order). ``a[ISAACLAB_TO_MUJOCO]``
# reorders an IsaacLab-ordered vector into MuJoCo order, and vice-versa.
ISAACLAB_TO_MUJOCO = np.array(
[
0,
3,
6,
9,
13,
17,
1,
4,
7,
10,
14,
18,
2,
5,
8,
11,
15,
19,
21,
23,
25,
27,
12,
16,
20,
22,
24,
26,
28,
],
dtype=np.int32,
)
# The two orderings are inverses of each other, so derive one from the other (argsort) to
# guarantee they can never drift out of sync.
MUJOCO_TO_ISAACLAB = np.argsort(ISAACLAB_TO_MUJOCO).astype(np.int32)
REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry") REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16)) REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
@@ -109,9 +68,8 @@ def make_locomotion_controller(name: str | None):
if name is None: if name is None:
return None return None
controllers = { controllers = {
"GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion", "GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
"HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion", "HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
"SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
} }
module_path = controllers.get(name) module_path = controllers.get(name)
if module_path is None: if module_path is None:
@@ -21,7 +21,7 @@ import numpy as np
import onnxruntime as ort import onnxruntime as ort
from huggingface_hub import hf_hub_download from huggingface_hub import hf_hub_download
from ..g1_utils import ( from .g1_utils import (
REMOTE_AXES, REMOTE_AXES,
REMOTE_BUTTONS, REMOTE_BUTTONS,
G1_29_JointIndex, G1_29_JointIndex,
@@ -22,7 +22,7 @@ import onnx
import onnxruntime as ort import onnxruntime as ort
from huggingface_hub import hf_hub_download from huggingface_hub import hf_hub_download
from ..g1_utils import ( from .g1_utils import (
REMOTE_AXES, REMOTE_AXES,
G1_29_JointArmIndex, G1_29_JointArmIndex,
G1_29_JointIndex, G1_29_JointIndex,
+40 -202
View File
@@ -34,6 +34,7 @@ from .config_unitree_g1 import UnitreeG1Config
from .g1_kinematics import G1_29_ArmIK from .g1_kinematics import G1_29_ArmIK
from .g1_utils import ( from .g1_utils import (
REMOTE_AXES, REMOTE_AXES,
REMOTE_KEYS,
G1_29_JointArmIndex, G1_29_JointArmIndex,
G1_29_JointIndex, G1_29_JointIndex,
default_remote_input, default_remote_input,
@@ -105,47 +106,6 @@ class G1_29_LowState: # noqa: N801
mode_machine: int = 0 # Robot mode mode_machine: int = 0 # Robot mode
def lowstate_to_obs(lowstate) -> dict:
"""Build a robot observation dict from a Unitree lowstate.
Shared by ``UnitreeG1.get_observation`` and the SONIC pipeline so the
lowstate -> obs mapping lives in exactly one place. Keys match the
``<joint>.q``/``imu.*`` schema consumed across the controllers.
"""
obs: dict = {}
for motor in G1_29_JointIndex:
idx = motor.value
obs[f"{motor.name}.q"] = lowstate.motor_state[idx].q
obs[f"{motor.name}.dq"] = lowstate.motor_state[idx].dq
obs[f"{motor.name}.tau"] = lowstate.motor_state[idx].tau_est
imu = lowstate.imu_state
if imu.gyroscope:
obs["imu.gyro.x"] = imu.gyroscope[0]
obs["imu.gyro.y"] = imu.gyroscope[1]
obs["imu.gyro.z"] = imu.gyroscope[2]
if imu.accelerometer:
obs["imu.accel.x"] = imu.accelerometer[0]
obs["imu.accel.y"] = imu.accelerometer[1]
obs["imu.accel.z"] = imu.accelerometer[2]
if imu.quaternion:
obs["imu.quat.w"] = imu.quaternion[0]
obs["imu.quat.x"] = imu.quaternion[1]
obs["imu.quat.y"] = imu.quaternion[2]
obs["imu.quat.z"] = imu.quaternion[3]
if imu.rpy:
obs["imu.rpy.roll"] = imu.rpy[0]
obs["imu.rpy.pitch"] = imu.rpy[1]
obs["imu.rpy.yaw"] = imu.rpy[2]
wr = getattr(lowstate, "wireless_remote", None)
if wr:
obs["wireless_remote"] = bytes(wr) if not isinstance(wr, (bytes, bytearray)) else wr
return obs
class UnitreeG1(Robot): class UnitreeG1(Robot):
config_class = UnitreeG1Config config_class = UnitreeG1Config
name = "unitree_g1" name = "unitree_g1"
@@ -188,60 +148,22 @@ class UnitreeG1(Robot):
self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
# Lower-body / whole-body controller loaded dynamically # Lower-body controller loaded dynamically
self.controller: LocomotionController | None = make_locomotion_controller(config.controller) self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
# A SONIC whole-body controller always runs in token mode: it holds a neutral
# token until the first real one arrives, then holds the last token between ticks.
if self.controller is not None and hasattr(self.controller, "token_mode"):
self.controller.token_mode = True
# Controller thread state # Controller thread state
self._controller_thread = None self._controller_thread = None
# When set, the controller loop stops publishing low commands so reset() can
# drive the joints directly without two publishers fighting (single-publisher).
self._controller_paused = threading.Event()
self._controller_action_lock = threading.Lock() self._controller_action_lock = threading.Lock()
self.controller_input = default_remote_input() self.controller_input = default_remote_input()
self.controller_output = {} self.controller_output = {}
# Token-mode state: last 64-D SONIC latent token commanded by the policy,
# echoed back as ``observation.state`` so a token-output VLA closes the loop
# on its own previous token. Implicit whenever the SONIC whole-body controller
# is active. Seeded to zeros; the controller's startup blend eases joints in.
self._last_token: np.ndarray | None = None
if self._sonic_token:
from .controllers.sonic_whole_body import TOKEN_DIM
self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32)
@property
def _sonic_token(self) -> bool:
"""Whether the SONIC whole-body decoder is active.
A SONIC controller consumes a 64-D latent motion token as its action and echoes
the last commanded token as ``observation.state``. Keyed purely off the selected
controller so the token interface is implicit -- no separate config flag.
"""
return self.config.controller == "SonicWholeBodyController"
def _subscribe_lowstate(self): # polls robot state @ 250Hz def _subscribe_lowstate(self): # polls robot state @ 250Hz
while not self._shutdown_event.is_set(): while not self._shutdown_event.is_set():
start_time = time.time() start_time = time.time()
# Step simulation if in simulation mode # Step simulation if in simulation mode
if self.config.is_simulation and self.sim_env is not None: if self.config.is_simulation and self.sim_env is not None:
try: self.sim_env.step()
self.sim_env.step()
except ValueError as e:
# Startup race: the sim thread can step once before reset() has
# written a valid base pose, giving a zero-norm pelvis quaternion
# (scipy>=1.11 raises instead of normalizing). Skip and retry so
# the thread survives instead of dying and freezing the sim.
if "zero norm" not in str(e).lower():
raise
time.sleep(self.control_dt)
continue
msg = self.lowstate_subscriber.Read() msg = self.lowstate_subscriber.Read()
if msg is not None: if msg is not None:
@@ -309,38 +231,15 @@ class UnitreeG1(Robot):
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1) features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
return features return features
@property
def _token_state_ft(self) -> dict[str, type]:
"""64-D SONIC latent-token proprio state (``motion_token_state.{i}.pos``).
Exposed only when a SONIC whole-body controller is active; aggregated by the
rollout into a 64-D ``observation.state`` (the last token the policy commanded).
"""
if not self._sonic_token:
return {}
from .controllers.sonic_whole_body import TOKEN_DIM, token_state_key
return {token_state_key(i): float for i in range(TOKEN_DIM)}
@cached_property @cached_property
def observation_features(self) -> dict[str, type | tuple]: def observation_features(self) -> dict[str, type | tuple]:
return {**self._motors_ft, **self._token_state_ft, **self._cameras_ft} return {**self._motors_ft, **self._cameras_ft}
@cached_property @cached_property
def action_features(self) -> dict[str, type]: def action_features(self) -> dict[str, type]:
# No controller configured at all: raw 29-DoF joint teleop.
if self.controller is None: if self.controller is None:
return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex} return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
# Token-output VLA (SONIC decoder): advertise a 64-D latent-token action space
# (``motion_token.{i}.pos``) so ``lerobot-rollout`` maps a 64-D policy output
# straight onto the decoder, bypassing the encoder.
if self._sonic_token:
from .controllers.sonic_whole_body import TOKEN_DIM, token_action_key
return {token_action_key(i): float for i in range(TOKEN_DIM)}
# Locomotion controllers (GR00T / Holosoma): arm joint targets + joystick axes.
arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex} arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
remote_features = dict.fromkeys(REMOTE_AXES, float) remote_features = dict.fromkeys(REMOTE_AXES, float)
return {**arm_features, **remote_features} return {**arm_features, **remote_features}
@@ -356,11 +255,6 @@ class UnitreeG1(Robot):
while not self._shutdown_event.is_set(): while not self._shutdown_event.is_set():
start_time = time.time() start_time = time.time()
# Paused during reset() so the reset routine is the sole low-cmd publisher.
if self._controller_paused.is_set():
time.sleep(control_dt)
continue
with self._lowstate_lock: with self._lowstate_lock:
lowstate = self._lowstate lowstate = self._lowstate
@@ -449,9 +343,6 @@ class UnitreeG1(Robot):
self.kp = np.array(self.config.kp, dtype=np.float32) self.kp = np.array(self.config.kp, dtype=np.float32)
self.kd = np.array(self.config.kd, dtype=np.float32) self.kd = np.array(self.config.kd, dtype=np.float32)
if self.controller is not None and hasattr(self.controller, "kp"):
self.kp = np.array(self.controller.kp, dtype=np.float32)
self.kd = np.array(self.controller.kd, dtype=np.float32)
for joint in G1_29_JointIndex: for joint in G1_29_JointIndex:
self.msg.motor_cmd[joint].mode = 1 self.msg.motor_cmd[joint].mode = 1
@@ -500,10 +391,6 @@ class UnitreeG1(Robot):
if self._controller_thread.is_alive(): if self._controller_thread.is_alive():
logger.warning("Controller thread did not stop cleanly") logger.warning("Controller thread did not stop cleanly")
# Release controller resources (e.g. SONIC decoder sessions).
if self.controller is not None and hasattr(self.controller, "shutdown"):
self.controller.shutdown()
# Close simulation environment # Close simulation environment
if self.config.is_simulation and self.sim_env is not None: if self.config.is_simulation and self.sim_env is not None:
try: try:
@@ -574,15 +461,6 @@ class UnitreeG1(Robot):
if lowstate.wireless_remote: if lowstate.wireless_remote:
obs["wireless_remote"] = lowstate.wireless_remote obs["wireless_remote"] = lowstate.wireless_remote
# Token mode: echo the last commanded latent token as observation.state so a
# token-output VLA closes the loop on its own previous token.
if self._sonic_token:
from .controllers.sonic_whole_body import token_state_key
token = self._last_token if self._last_token is not None else []
for i, v in enumerate(token):
obs[token_state_key(i)] = float(v)
# Cameras - read images from ZMQ cameras # Cameras - read images from ZMQ cameras
for cam_name, cam in self._cameras.items(): for cam_name, cam in self._cameras.items():
if getattr(cam, "use_rgb", True): if getattr(cam, "use_rgb", True):
@@ -595,22 +473,9 @@ class UnitreeG1(Robot):
def send_action(self, action: RobotAction) -> RobotAction: def send_action(self, action: RobotAction) -> RobotAction:
action_to_publish = action action_to_publish = action
if self.controller is not None: if self.controller is not None:
# SONIC decoder: pull the 64-D latent token out of the action and remember it
# for the observation.state echo. The controller thread reads it back from
# controller_input (populated below) and decodes it into a 29-DoF command.
if self._sonic_token:
from .controllers.sonic_whole_body import _extract_token_from_action
token = _extract_token_from_action(action)
if token is not None:
self._last_token = token
self._update_controller_action(action)
# Full-body controllers (SONIC) own the whole 29-DoF command; nothing to
# publish here (the controller thread is the sole publisher).
if getattr(self.controller, "full_body", False):
return action
# Controller thread owns legs/waist. Here we only update joystick inputs # Controller thread owns legs/waist. Here we only update joystick inputs
# and publish arm targets from the teleoperator. # and publish arm targets from the teleoperator.
self._update_controller_action(action)
arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex) arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
action_to_publish = { action_to_publish = {
key: value key: value
@@ -638,17 +503,11 @@ class UnitreeG1(Robot):
return action return action
def _update_controller_action(self, action: RobotAction) -> None: def _update_controller_action(self, action: RobotAction) -> None:
"""Update controller input state from an incoming teleop action. """Update controller input state from incoming teleop action."""
Controller-agnostic: every value-carrying key (locomotion ``remote.*`` axes or
SONIC ``motion_token.*`` values) is forwarded verbatim into ``controller_input``
and each controller extracts only the keys it understands. The robot deliberately
does not enumerate any controller's key schema here.
"""
with self._controller_action_lock: with self._controller_action_lock:
for key, value in action.items(): for key in REMOTE_KEYS:
if isinstance(key, str) and value is not None: if key in action:
self.controller_input[key] = value self.controller_input[key] = action[key]
@property @property
def is_calibrated(self) -> bool: def is_calibrated(self) -> bool:
@@ -678,64 +537,43 @@ class UnitreeG1(Robot):
if default_positions is None: if default_positions is None:
default_positions = np.array(self.config.default_positions, dtype=np.float32) default_positions = np.array(self.config.default_positions, dtype=np.float32)
# Full-body controllers (SONIC) own the whole 29-DoF command and ignore if self.config.is_simulation and self.sim_env is not None:
# ``<joint>.q`` in send_action(), so reset() must publish the default pose self.sim_env.reset()
# directly. Pause the background controller first so the two aren't both writing self.publish_lowcmd(
# low commands while the robot moves to the default pose. {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
full_body = getattr(self.controller, "full_body", False) )
paused = False else:
if full_body and self._controller_thread is not None: total_time = 3.0
self._controller_paused.set() num_steps = int(total_time / control_dt)
paused = True
time.sleep(control_dt) # let any in-flight controller tick settle
try: # get current state
if self.config.is_simulation and self.sim_env is not None: obs = self.get_observation()
self.sim_env.reset()
self.publish_lowcmd(
{f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
)
else:
total_time = 3.0
num_steps = int(total_time / control_dt)
# get current state # record current positions
obs = self.get_observation() init_dof_pos = np.zeros(29, dtype=np.float32)
for motor in G1_29_JointIndex:
init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
# record current positions # Interpolate to default position
init_dof_pos = np.zeros(29, dtype=np.float32) for step in range(num_steps):
start_time = time.time()
alpha = step / num_steps
action_dict = {}
for motor in G1_29_JointIndex: for motor in G1_29_JointIndex:
init_dof_pos[motor.value] = obs[f"{motor.name}.q"] target_pos = default_positions[motor.value]
interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
action_dict[f"{motor.name}.q"] = float(interp_pos)
# Interpolate to default position self.send_action(action_dict)
for step in range(num_steps):
start_time = time.time()
alpha = step / num_steps # Maintain constant control rate
action_dict = {} elapsed = time.time() - start_time
for motor in G1_29_JointIndex: sleep_time = max(0, control_dt - elapsed)
target_pos = default_positions[motor.value] time.sleep(sleep_time)
interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
action_dict[f"{motor.name}.q"] = float(interp_pos)
# Full-body controllers no-op in send_action(); publish the pose # Reset controller internal state (gait phase, obs history, etc.)
# directly (arm-only controllers keep the send_action() path). if self.controller is not None and hasattr(self.controller, "reset"):
if full_body: self.controller.reset()
self.publish_lowcmd(action_dict)
else:
self.send_action(action_dict)
# Maintain constant control rate
elapsed = time.time() - start_time
sleep_time = max(0, control_dt - elapsed)
time.sleep(sleep_time)
# Reset controller internal state (gait phase, obs history, etc.) before
# resuming so its buffers reflect the post-reset pose.
if self.controller is not None and hasattr(self.controller, "reset"):
self.controller.reset()
finally:
if paused:
self._controller_paused.clear()
logger.info("Reached default position") logger.info("Reached default position")
+9 -22
View File
@@ -24,7 +24,6 @@ from __future__ import annotations
import logging import logging
from dataclasses import dataclass, field from dataclasses import dataclass, field
from threading import Event from threading import Event
from typing import TYPE_CHECKING
import torch import torch
@@ -48,7 +47,6 @@ from lerobot.processor.relative_action_processor import RelativeActionsProcessor
from lerobot.robots import make_robot_from_config from lerobot.robots import make_robot_from_config
from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
from lerobot.utils.import_utils import _peft_available, require_package
from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
from .inference import ( from .inference import (
@@ -59,12 +57,6 @@ from .inference import (
) )
from .robot_wrapper import ThreadSafeRobot from .robot_wrapper import ThreadSafeRobot
if TYPE_CHECKING or _peft_available:
from peft import PeftConfig, PeftModel
else:
PeftConfig = None
PeftModel = None
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -169,7 +161,12 @@ class RolloutContext:
def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy: def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy:
"""Load policy weights, keeping adapter and base-model revisions independent.""" """Load policy weights, keeping adapter and base-model revisions independent."""
pretrained_revision = policy_config.pretrained_revision revision_resolver = getattr(policy_config, "get_hub_revision", None)
pretrained_revision = (
revision_resolver(policy_config.pretrained_path, policy_config.pretrained_revision)
if callable(revision_resolver)
else policy_config.pretrained_revision
)
policy_class = get_policy_class(policy_config.type) policy_class = get_policy_class(policy_config.type)
if not policy_config.use_peft: if not policy_config.use_peft:
@@ -179,7 +176,7 @@ def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy
revision=pretrained_revision, revision=pretrained_revision,
) )
require_package("peft", extra="peft") from peft import PeftConfig, PeftModel
peft_path = policy_config.pretrained_path peft_path = policy_config.pretrained_path
peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision) peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
@@ -302,22 +299,12 @@ def build_rollout_context(
# ``observation_features`` values are either a tuple (camera shape) or the # ``observation_features`` values are either a tuple (camera shape) or the
# ``float`` type itself used as a sentinel for scalar motor features — # ``float`` type itself used as a sentinel for scalar motor features —
# see ``dict[str, type | tuple]`` annotation on ``Robot.observation_features``. # see ``dict[str, type | tuple]`` annotation on ``Robot.observation_features``.
# Keep cameras (tuple) plus both joint-position (.pos) and base-velocity (.vel)
# scalar state features. LeKiwi's observation.state is 9-dim (6 arm .pos +
# x/y/theta.vel) and the policy was trained/normalized on all 9; the old .pos-only
# filter fed a 6-dim state into a 9-dim normalizer → RuntimeError (size 6 vs 9).
# Pure-arm robots have no .vel state keys, so this is a no-op for them.
observation_features_hw = { observation_features_hw = {
k: v k: v
for k, v in all_obs_features.items() for k, v in all_obs_features.items()
if isinstance(v, tuple) or (v is float and k.endswith((".pos", ".vel"))) if isinstance(v, tuple) or (v is float and k.endswith(".pos"))
} }
# Keep both joint-position (.pos) and base-velocity (.vel) action features so action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith(".pos")}
# mobile manipulators command the base too (e.g. LeKiwi: 6 arm .pos +
# x/y/theta.vel = 9-dim action). Pure-arm robots have no .vel keys, so this is
# a no-op for them. Without the .vel keys the base velocities are silently
# dropped from dataset_features[ACTION]/ordered_action_keys and the base never moves.
action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith((".pos", ".vel"))}
# The action side is always needed: sync inference reads action names from # The action side is always needed: sync inference reads action names from
# ``dataset_features[ACTION]`` to map policy tensors back to robot actions. # ``dataset_features[ACTION]`` to map policy tensors back to robot actions.
@@ -36,7 +36,6 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
import argparse import argparse
import concurrent.futures import concurrent.futures
import logging import logging
import os
from pathlib import Path from pathlib import Path
import numpy as np import numpy as np
@@ -53,7 +52,6 @@ from lerobot.datasets import (
get_feature_stats, get_feature_stats,
write_stats, write_stats,
) )
from lerobot.datasets.compute_stats import sample_indices
from lerobot.utils.utils import init_logging from lerobot.utils.utils import init_logging
@@ -79,14 +77,12 @@ def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[s
return False return False
def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampling: bool = True) -> dict: def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
"""Process a single episode and return its statistics. """Process a single episode and return its statistics.
Args: Args:
dataset: The LeRobot dataset dataset: The LeRobot dataset
episode_idx: Index of the episode to process episode_idx: Index of the episode to process
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
Returns: Returns:
Dictionary containing episode statistics Dictionary containing episode statistics
@@ -96,31 +92,16 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampli
start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"] start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"] end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]
episode_len = end_idx - start_idx
# Images/video are the memory hog, so sub-sample those frames per episode;
# numeric columns are cheap, so read them in full (exact).
image_keys = [k for k in dataset.features if dataset.features[k]["dtype"] in ("image", "video")]
numeric_keys = [
k for k in dataset.features if dataset.features[k]["dtype"] not in ("image", "video", "string")
]
collected_data: dict[str, list] = {} collected_data: dict[str, list] = {}
for idx in range(start_idx, end_idx):
item = dataset[idx]
for key, value in item.items():
if key not in dataset.features:
continue
# Numeric features: every frame, read directly from the underlying table. if key not in collected_data:
if numeric_keys: collected_data[key] = []
numeric_cols = dataset.hf_dataset.select_columns(numeric_keys)[start_idx:end_idx] collected_data[key].append(value)
for key in numeric_keys:
collected_data[key] = [torch.as_tensor(v) for v in numeric_cols[key]]
# Image/video features: decode only a sampled subset of frames.
if image_keys:
sampled_offsets = sample_indices(episode_len) if use_sampling else list(range(episode_len))
for offset in sampled_offsets:
item = dataset[start_idx + offset]
for key in image_keys:
if key in item:
collected_data.setdefault(key, []).append(item[key])
ep_stats = {} ep_stats = {}
for key, data_list in collected_data.items(): for key, data_list in collected_data.items():
@@ -150,13 +131,11 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampli
return ep_stats return ep_stats
def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]: def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dict]:
"""Compute quantile statistics for all episodes in the dataset. """Compute quantile statistics for all episodes in the dataset.
Args: Args:
dataset: The LeRobot dataset to compute statistics for dataset: The LeRobot dataset to compute statistics for
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
Returns: Returns:
Dictionary containing aggregated statistics with quantiles Dictionary containing aggregated statistics with quantiles
@@ -174,15 +153,15 @@ def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bo
if has_videos: if has_videos:
logging.info("Dataset contains video keys - using sequential processing for thread safety") logging.info("Dataset contains video keys - using sequential processing for thread safety")
for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"): for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
ep_stats = process_single_episode(dataset, episode_idx, use_sampling) ep_stats = process_single_episode(dataset, episode_idx)
episode_stats_list.append(ep_stats) episode_stats_list.append(ep_stats)
else: else:
logging.info("Dataset has no video keys - using parallel processing for better performance") logging.info("Dataset has no video keys - using parallel processing for better performance")
max_workers = min(dataset.num_episodes, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16))) max_workers = min(dataset.num_episodes, 16)
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
future_to_episode = { future_to_episode = {
executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx executor.submit(process_single_episode, dataset, episode_idx): episode_idx
for episode_idx in range(dataset.num_episodes) for episode_idx in range(dataset.num_episodes)
} }
@@ -209,7 +188,6 @@ def augment_dataset_with_quantile_stats(
repo_id: str, repo_id: str,
root: str | Path | None = None, root: str | Path | None = None,
overwrite: bool = False, overwrite: bool = False,
use_sampling: bool = True,
) -> None: ) -> None:
"""Augment a dataset with quantile statistics if they are missing. """Augment a dataset with quantile statistics if they are missing.
@@ -217,8 +195,6 @@ def augment_dataset_with_quantile_stats(
repo_id: Repository ID of the dataset repo_id: Repository ID of the dataset
root: Local root directory for the dataset root: Local root directory for the dataset
overwrite: Overwrite existing quantile statistics if they already exist overwrite: Overwrite existing quantile statistics if they already exist
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
""" """
logging.info(f"Loading dataset: {repo_id}") logging.info(f"Loading dataset: {repo_id}")
dataset = LeRobotDataset( dataset = LeRobotDataset(
@@ -232,7 +208,7 @@ def augment_dataset_with_quantile_stats(
logging.info("Dataset does not contain quantile statistics. Computing them now...") logging.info("Dataset does not contain quantile statistics. Computing them now...")
new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling) new_stats = compute_quantile_stats_for_dataset(dataset)
logging.info("Updating dataset metadata with new quantile statistics") logging.info("Updating dataset metadata with new quantile statistics")
dataset.meta.stats = new_stats dataset.meta.stats = new_stats
@@ -272,14 +248,6 @@ def main():
action="store_true", action="store_true",
help="Overwrite existing quantile statistics if they already exist", help="Overwrite existing quantile statistics if they already exist",
) )
parser.add_argument(
"--no-sampling",
action="store_true",
help=(
"Compute stats over every frame (exact, higher memory). By default, "
"image/video frames are sub-sampled per episode to bound memory."
),
)
args = parser.parse_args() args = parser.parse_args()
root = Path(args.root) if args.root else None root = Path(args.root) if args.root else None
@@ -290,7 +258,6 @@ def main():
repo_id=args.repo_id, repo_id=args.repo_id,
root=root, root=root,
overwrite=args.overwrite, overwrite=args.overwrite,
use_sampling=not args.no_sampling,
) )
@@ -94,8 +94,6 @@ from lerobot.datasets.video_utils import concatenate_video_files, get_video_dura
from lerobot.utils.constants import HF_LEROBOT_HOME from lerobot.utils.constants import HF_LEROBOT_HOME
from lerobot.utils.utils import flatten_dict, init_logging from lerobot.utils.utils import flatten_dict, init_logging
logger = logging.getLogger(__name__)
V21 = "v2.1" V21 = "v2.1"
V30 = "v3.0" V30 = "v3.0"
@@ -478,11 +476,11 @@ def convert_dataset(
# First check if the dataset already has a v3.0 version # First check if the dataset already has a v3.0 version
if root is None and not force_conversion: if root is None and not force_conversion:
try: try:
logger.info("Trying to download v3.0 version of the dataset from the hub...") print("Trying to download v3.0 version of the dataset from the hub...")
snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id) snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id)
return return
except Exception: except Exception:
logger.info("Dataset does not have an uploaded v3.0 version. Continuing with conversion.") print("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
# Set root based on whether local dataset path is provided # Set root based on whether local dataset path is provided
use_local_dataset = False use_local_dataset = False
@@ -490,7 +488,7 @@ def convert_dataset(
if root.exists(): if root.exists():
validate_local_dataset_version(root) validate_local_dataset_version(root)
use_local_dataset = True use_local_dataset = True
logger.info(f"Using local dataset at {root}") print(f"Using local dataset at {root}")
old_root = root.parent / f"{root.name}_old" old_root = root.parent / f"{root.name}_old"
new_root = root.parent / f"{root.name}_v30" new_root = root.parent / f"{root.name}_v30"
@@ -525,7 +523,7 @@ def convert_dataset(
try: try:
hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset") hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
except (HTTPError, RevisionNotFoundError) as e: except (HTTPError, RevisionNotFoundError) as e:
logger.warning(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})") print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
pass pass
hub_api.delete_files( hub_api.delete_files(
delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"], delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"],
+7 -6
View File
@@ -154,14 +154,14 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
repo_id = cfg.new_repo_id or cfg.repo_id repo_id = cfg.new_repo_id or cfg.repo_id
commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)" commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
api = HfApi() api = HfApi()
logger.info(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...") print(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...", flush=True)
api.create_repo( api.create_repo(
repo_id=repo_id, repo_id=repo_id,
repo_type="dataset", repo_type="dataset",
private=cfg.push_private, private=cfg.push_private,
exist_ok=True, exist_ok=True,
) )
logger.info(f"[lerobot-annotate] uploading {root} -> {repo_id}...") print(f"[lerobot-annotate] uploading {root} -> {repo_id}...", flush=True)
commit_info = api.upload_folder( commit_info = api.upload_folder(
folder_path=str(root), folder_path=str(root),
repo_id=repo_id, repo_id=repo_id,
@@ -172,7 +172,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
# at the source dataset; a fresh card is generated below instead. # at the source dataset; a fresh card is generated below instead.
ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"], ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
) )
logger.info(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}") print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True)
dataset_info = load_info(root) dataset_info = load_info(root)
card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id) card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
@@ -200,13 +200,14 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
with suppress(RevisionNotFoundError): with suppress(RevisionNotFoundError):
api.delete_tag(repo_id, tag=version_tag, repo_type="dataset") api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
api.create_tag(**tag_kwargs) api.create_tag(**tag_kwargs)
logger.info(f"[lerobot-annotate] tagged {repo_id} as {version_tag}") print(f"[lerobot-annotate] tagged {repo_id} as {version_tag}", flush=True)
except Exception as exc: # noqa: BLE001 except Exception as exc: # noqa: BLE001
logger.warning( print(
f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. " f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. "
"Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. " "Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. "
"Run: from huggingface_hub import HfApi; " "Run: from huggingface_hub import HfApi; "
f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)" f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)",
flush=True,
) )
+1 -3
View File
@@ -89,8 +89,6 @@ from lerobot.datasets import LeRobotDataset
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
from lerobot.utils.utils import init_logging from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
DEFAULT_FOXGLOVE_PORT = 8765 DEFAULT_FOXGLOVE_PORT = 8765
DEFAULT_RERUN_PORT = 9090 DEFAULT_RERUN_PORT = 9090
@@ -301,7 +299,7 @@ def visualize_dataset(
while True: while True:
time.sleep(1) time.sleep(1)
except KeyboardInterrupt: except KeyboardInterrupt:
logger.info("Ctrl-C received. Exiting.") print("Ctrl-C received. Exiting.")
def main(): def main():
+14 -20
View File
@@ -62,7 +62,7 @@ from dataclasses import asdict
from functools import partial from functools import partial
from pathlib import Path from pathlib import Path
from pprint import pformat from pprint import pformat
from typing import TYPE_CHECKING, Any, TypedDict from typing import Any, TypedDict
import einops import einops
import gymnasium as gym import gymnasium as gym
@@ -87,7 +87,7 @@ from lerobot.processor import PolicyProcessorPipeline
from lerobot.types import PolicyAction from lerobot.types import PolicyAction
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
from lerobot.utils.device_utils import get_safe_torch_device from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package from lerobot.utils.import_utils import register_third_party_plugins
from lerobot.utils.io_utils import write_video from lerobot.utils.io_utils import write_video
from lerobot.utils.random_utils import set_seed from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import ( from lerobot.utils.utils import (
@@ -95,14 +95,6 @@ from lerobot.utils.utils import (
inside_slurm, inside_slurm,
) )
if TYPE_CHECKING or _peft_available:
from peft import PeftModel
else:
PeftModel = None
logger = logging.getLogger(__name__)
def _env_features_to_dataset_features(env_features: dict) -> dict: def _env_features_to_dataset_features(env_features: dict) -> dict:
"""Convert EnvConfig.features to the dict format expected by LeRobotDataset.create().""" """Convert EnvConfig.features to the dict format expected by LeRobotDataset.create()."""
@@ -452,11 +444,13 @@ def eval_policy(
exc = ValueError( exc = ValueError(
f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided." f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
) )
if not _peft_available: try:
raise exc from peft import PeftModel
require_package("peft", extra="peft")
if not isinstance(policy, PeftModel): if not isinstance(policy, PeftModel):
raise exc raise exc
except ImportError:
raise exc from None
start = time.time() start = time.time()
# Preserve the mode for direct callers. eval_policy_all scopes the mode # Preserve the mode for direct callers. eval_policy_all scopes the mode
@@ -564,7 +558,7 @@ def eval_policy(
if seeds: if seeds:
all_seeds.extend(seeds) all_seeds.extend(seeds)
else: else:
all_seeds.extend([None] * env.num_envs) all_seeds.append(None)
# FIXME: episode_data is either None or it doesn't exist # FIXME: episode_data is either None or it doesn't exist
if return_episode_data: if return_episode_data:
@@ -802,13 +796,13 @@ def eval_main(cfg: EvalPipelineConfig):
recording_repo_id=cfg.eval.recording_repo_id, recording_repo_id=cfg.eval.recording_repo_id,
recording_private=cfg.eval.recording_private, recording_private=cfg.eval.recording_private,
) )
logger.info("Overall Aggregated Metrics:") print("Overall Aggregated Metrics:")
logger.info(info["overall"]) print(info["overall"])
# Print per-suite stats # Print per-suite stats
for task_group, task_group_info in info.items(): for task_group, task_group_info in info.items():
logger.info(f"\nAggregated Metrics for {task_group}:") print(f"\nAggregated Metrics for {task_group}:")
logger.info(task_group_info) print(task_group_info)
# Close all vec envs # Close all vec envs
close_envs(envs) close_envs(envs)
+49 -41
View File
@@ -28,6 +28,7 @@ lerobot-find-cameras
# NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful. # NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
import argparse import argparse
import concurrent.futures
import logging import logging
import time import time
from pathlib import Path from pathlib import Path
@@ -39,7 +40,6 @@ from PIL import Image
from lerobot.cameras import ColorMode from lerobot.cameras import ColorMode
from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -132,7 +132,7 @@ def save_image(
camera_identifier: str | int, camera_identifier: str | int,
images_dir: Path, images_dir: Path,
camera_type: str, camera_type: str,
) -> None: ):
""" """
Saves a single image to disk using Pillow. Handles color conversion if necessary. Saves a single image to disk using Pillow. Handles color conversion if necessary.
""" """
@@ -151,7 +151,7 @@ def save_image(
logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}") logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}")
def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> dict[str, Any] | None: def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
"""Create and connect to a camera instance based on metadata.""" """Create and connect to a camera instance based on metadata."""
cam_type = cam_meta.get("type") cam_type = cam_meta.get("type")
cam_id = cam_meta.get("id") cam_id = cam_meta.get("id")
@@ -164,14 +164,12 @@ def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> di
cv_config = OpenCVCameraConfig( cv_config = OpenCVCameraConfig(
index_or_path=cam_id, index_or_path=cam_id,
color_mode=ColorMode.RGB, color_mode=ColorMode.RGB,
warmup_s=warmup_s,
) )
instance = OpenCVCamera(cv_config) instance = OpenCVCamera(cv_config)
elif cam_type == "RealSense": elif cam_type == "RealSense":
rs_config = RealSenseCameraConfig( rs_config = RealSenseCameraConfig(
serial_number_or_name=cam_id, serial_number_or_name=cam_id,
color_mode=ColorMode.RGB, color_mode=ColorMode.RGB,
warmup_s=warmup_s,
) )
instance = RealSenseCamera(rs_config) instance = RealSenseCamera(rs_config)
else: else:
@@ -189,7 +187,9 @@ def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> di
return None return None
def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_time: float) -> None: def process_camera_image(
cam_dict: dict[str, Any], output_dir: Path, current_time: float
) -> concurrent.futures.Future | None:
"""Capture and process an image from a single camera.""" """Capture and process an image from a single camera."""
cam = cam_dict["instance"] cam = cam_dict["instance"]
meta = cam_dict["meta"] meta = cam_dict["meta"]
@@ -199,7 +199,7 @@ def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_tim
try: try:
image_data = cam.read() image_data = cam.read()
save_image( return save_image(
image_data, image_data,
cam_id_str, cam_id_str,
output_dir, output_dir,
@@ -214,21 +214,21 @@ def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_tim
return None return None
def cleanup_camera(cam_dict: dict[str, Any]) -> None: def cleanup_cameras(cameras_to_use: list[dict[str, Any]]):
"""Disconnect all cameras.""" """Disconnect all cameras."""
logger.info(f"Disconnecting camera with ID {cam_dict['meta'].get('id')}...") logger.info(f"Disconnecting {len(cameras_to_use)} cameras...")
try: for cam_dict in cameras_to_use:
if cam_dict["instance"] and cam_dict["instance"].is_connected: try:
cam_dict["instance"].disconnect() if cam_dict["instance"] and cam_dict["instance"].is_connected:
except Exception as e: cam_dict["instance"].disconnect()
logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}") except Exception as e:
logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
def save_images_from_all_cameras( def save_images_from_all_cameras(
output_dir: Path, output_dir: Path,
record_time_s: float = 2.0, record_time_s: float = 2.0,
camera_type: str | None = None, camera_type: str | None = None,
warmup_s: int = 1,
): ):
""" """
Connects to detected cameras (optionally filtered by type) and saves images from each. Connects to detected cameras (optionally filtered by type) and saves images from each.
@@ -239,7 +239,6 @@ def save_images_from_all_cameras(
record_time_s: Duration in seconds to record images. record_time_s: Duration in seconds to record images.
camera_type: Optional string to filter cameras ("realsense" or "opencv"). camera_type: Optional string to filter cameras ("realsense" or "opencv").
If None, uses all detected cameras. If None, uses all detected cameras.
warmup_s: Duration in seconds to warmup camera before recording images.
""" """
output_dir.mkdir(parents=True, exist_ok=True) output_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Saving images to {output_dir}") logger.info(f"Saving images to {output_dir}")
@@ -249,32 +248,47 @@ def save_images_from_all_cameras(
logger.warning("No cameras detected matching the criteria. Cannot save images.") logger.warning("No cameras detected matching the criteria. Cannot save images.")
return return
logger.info( cameras_to_use = []
f"Starting image capture for {record_time_s} seconds from {len(all_camera_metadata)} cameras." for cam_meta in all_camera_metadata:
) camera_instance = create_camera_instance(cam_meta)
if camera_instance:
cameras_to_use.append(camera_instance)
try: if not cameras_to_use:
for cam_meta in all_camera_metadata: logger.warning("No cameras could be connected. Aborting image save.")
cam_dict = create_camera_instance(cam_meta, warmup_s=warmup_s) return
if cam_dict is None:
continue logger.info(f"Starting image capture for {record_time_s} seconds from {len(cameras_to_use)} cameras.")
start_time = time.perf_counter() start_time = time.perf_counter()
with concurrent.futures.ThreadPoolExecutor(max_workers=len(cameras_to_use) * 2) as executor:
try:
while time.perf_counter() - start_time < record_time_s: while time.perf_counter() - start_time < record_time_s:
futures = []
current_capture_time = time.perf_counter() current_capture_time = time.perf_counter()
process_camera_image(cam_dict, output_dir, current_capture_time)
cleanup_camera(cam_dict) for cam_dict in cameras_to_use:
except KeyboardInterrupt: future = process_camera_image(cam_dict, output_dir, current_capture_time)
logger.info("Capture interrupted by user.") if future:
finally: futures.append(future)
print(f"Image capture finished. Images saved to {output_dir}")
if futures:
concurrent.futures.wait(futures)
except KeyboardInterrupt:
logger.info("Capture interrupted by user.")
finally:
print("\nFinalizing image saving...")
executor.shutdown(wait=True)
cleanup_cameras(cameras_to_use)
print(f"Image capture finished. Images saved to {output_dir}")
def main(): def main():
init_logging()
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description="Unified camera utility script for listing cameras and capturing images." description="Unified camera utility script for listing cameras and capturing images."
) )
parser.add_argument( parser.add_argument(
"camera_type", "camera_type",
type=str, type=str,
@@ -292,14 +306,8 @@ def main():
parser.add_argument( parser.add_argument(
"--record-time-s", "--record-time-s",
type=float, type=float,
default=2.0, default=6.0,
help="Time duration to attempt capturing frames. Default: 2 seconds.", help="Time duration to attempt capturing frames. Default: 6 seconds.",
)
parser.add_argument(
"--warmup-s",
type=int,
default=1,
help="Time duration to warmup camera before attempting to capture frames. Default: 1 second.",
) )
args = parser.parse_args() args = parser.parse_args()
save_images_from_all_cameras(**vars(args)) save_images_from_all_cameras(**vars(args))
-1
View File
@@ -165,7 +165,6 @@ from lerobot.robots import ( # noqa: F401
earthrover_mini_plus, earthrover_mini_plus,
hope_jr, hope_jr,
koch_follower, koch_follower,
lekiwi,
omx_follower, omx_follower,
openarm_follower, openarm_follower,
reachy2, reachy2,
+17 -24
View File
@@ -22,8 +22,7 @@ import dataclasses
import logging import logging
import sys import sys
import time import time
from collections.abc import Iterator from contextlib import nullcontext
from contextlib import contextmanager, nullcontext
from pprint import pformat from pprint import pformat
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
@@ -58,7 +57,7 @@ from lerobot.optim.factory import make_optimizer_and_scheduler
from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
from lerobot.rewards import make_reward_pre_post_processors from lerobot.rewards import make_reward_pre_post_processors
from lerobot.utils.collate import lerobot_collate_fn from lerobot.utils.collate import lerobot_collate_fn
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package from lerobot.utils.import_utils import register_third_party_plugins
from lerobot.utils.logging_utils import AverageMeter, MetricsTracker from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
from lerobot.utils.random_utils import set_seed from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import ( from lerobot.utils.utils import (
@@ -69,28 +68,9 @@ from lerobot.utils.utils import (
inside_slurm, inside_slurm,
) )
if TYPE_CHECKING or _peft_available:
from peft import PeftModel
else:
PeftModel = None
from .lerobot_eval import eval_policy_all from .lerobot_eval import eval_policy_all
@contextmanager
def _make_eval_envs(cfg: TrainPipelineConfig) -> Iterator[dict[str, dict[int, Any]]]:
"""Create evaluation environments for one run and always dispose of them."""
envs = make_env(
cfg.env,
n_envs=cfg.eval.batch_size,
use_async_envs=cfg.eval.use_async_envs,
)
try:
yield envs
finally:
close_envs(envs)
def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]: def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
"""Return worker-only DataLoader options, disabling them for single-process loading.""" """Return worker-only DataLoader options, disabling them for single-process loading."""
workers_enabled = cfg.num_workers > 0 workers_enabled = cfg.num_workers > 0
@@ -227,6 +207,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.job.is_remote: if cfg.job.is_remote:
return submit_to_hf(cfg) return submit_to_hf(cfg)
from lerobot.utils.import_utils import require_package
require_package("accelerate", extra="training") require_package("accelerate", extra="training")
from accelerate import Accelerator from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs, DistributedType from accelerate.utils import DistributedDataParallelKwargs, DistributedType
@@ -295,6 +277,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if not is_main_process: if not is_main_process:
dataset, eval_dataset = make_train_eval_datasets(cfg) dataset, eval_dataset = make_train_eval_datasets(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,
# using the eval.py instead, with gym_dora environment and dora-rs.
eval_env = None
if cfg.env_eval_freq > 0 and cfg.env is not None and is_main_process:
logging.info("Creating env")
eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
if cfg.is_reward_model_training: if cfg.is_reward_model_training:
if is_main_process: if is_main_process:
logging.info("Creating reward model") logging.info("Creating reward model")
@@ -322,7 +312,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.peft is not None: if cfg.peft is not None:
if cfg.is_reward_model_training: if cfg.is_reward_model_training:
raise ValueError("PEFT is only supported for policy training. ") raise ValueError("PEFT is only supported for policy training. ")
require_package("peft", extra="peft") from peft import PeftModel
if isinstance(policy, PeftModel): if isinstance(policy, PeftModel):
logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.") logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.")
@@ -702,7 +692,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process: if is_main_process:
step_id = get_step_identifier(step, cfg.steps) step_id = get_step_identifier(step, cfg.steps)
logging.info(f"Eval policy at step {step}") logging.info(f"Eval policy at step {step}")
with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast(): with torch.no_grad(), accelerator.autocast():
eval_info = eval_policy_all( eval_info = eval_policy_all(
envs=eval_env, # dict[suite][task_id] -> vec_env envs=eval_env, # dict[suite][task_id] -> vec_env
policy=accelerator.unwrap_model(policy), policy=accelerator.unwrap_model(policy),
@@ -750,6 +740,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process: if is_main_process:
progbar.close() progbar.close()
if eval_env:
close_envs(eval_env)
is_fsdp = accelerator.distributed_type == DistributedType.FSDP is_fsdp = accelerator.distributed_type == DistributedType.FSDP
model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None
if is_main_process: if is_main_process:
+56 -58
View File
@@ -45,7 +45,6 @@ lerobot-train-tokenizer \
""" """
import json import json
import logging
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
@@ -64,9 +63,6 @@ else:
from lerobot.configs import NormalizationMode, parser from lerobot.configs import NormalizationMode, parser
from lerobot.datasets import LeRobotDataset from lerobot.datasets import LeRobotDataset
from lerobot.utils.constants import ACTION, OBS_STATE from lerobot.utils.constants import ACTION, OBS_STATE
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
@dataclass @dataclass
@@ -278,8 +274,11 @@ def process_episode(args):
return action_chunks return action_chunks
except Exception: except Exception as e:
logger.exception("Error processing episode %s", ep_idx) print(f"Error processing episode {ep_idx}: {e}")
import traceback
traceback.print_exc()
return None return None
@@ -301,10 +300,10 @@ def train_fast_tokenizer(
Returns: Returns:
Trained FAST tokenizer Trained FAST tokenizer
""" """
logger.info(f"Training FAST tokenizer on {len(action_chunks)} action chunks...") print(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
logger.info(f"Action chunk shape: {action_chunks.shape}") print(f"Action chunk shape: {action_chunks.shape}")
logger.info(f"Vocab size: {vocab_size}") print(f"Vocab size: {vocab_size}")
logger.info(f"DCT scale: {scale}") print(f"DCT scale: {scale}")
# download the tokenizer source code (not pretrained weights) # download the tokenizer source code (not pretrained weights)
# we'll train a new tokenizer on our own data # we'll train a new tokenizer on our own data
@@ -315,7 +314,7 @@ def train_fast_tokenizer(
# train the new tokenizer on our action data using .fit() # train the new tokenizer on our action data using .fit()
# this trains the BPE tokenizer on DCT coefficients # this trains the BPE tokenizer on DCT coefficients
logger.info("Training new tokenizer (this may take a few minutes)...") print("Training new tokenizer (this may take a few minutes)...")
tokenizer = base_tokenizer.fit( tokenizer = base_tokenizer.fit(
action_data_list, action_data_list,
scale=scale, scale=scale,
@@ -323,21 +322,21 @@ def train_fast_tokenizer(
time_horizon=action_chunks.shape[1], # action_horizon time_horizon=action_chunks.shape[1], # action_horizon
action_dim=action_chunks.shape[2], # encoded dimensions action_dim=action_chunks.shape[2], # encoded dimensions
) )
logger.info("✓ Tokenizer training complete!") print("✓ Tokenizer training complete!")
# validate it works # validate it works
sample_chunk = action_chunks[0] sample_chunk = action_chunks[0]
encoded = tokenizer(sample_chunk[None])[0] encoded = tokenizer(sample_chunk[None])[0]
if isinstance(encoded, list): if isinstance(encoded, list):
encoded = np.array(encoded) encoded = np.array(encoded)
logger.info(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}") print(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
return tokenizer return tokenizer
def compute_compression_stats(tokenizer, action_chunks: np.ndarray): def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
"""Compute compression statistics.""" """Compute compression statistics."""
logger.info("\nComputing compression statistics...") print("\nComputing compression statistics...")
# sample for stats (use max 1000 chunks for speed) # sample for stats (use max 1000 chunks for speed)
sample_size = min(1000, len(action_chunks)) sample_size = min(1000, len(action_chunks))
@@ -367,12 +366,12 @@ def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
"max_token_length": float(np.max(token_lengths)), "max_token_length": float(np.max(token_lengths)),
} }
logger.info("Compression Statistics:") print("Compression Statistics:")
logger.info(f" Average compression ratio: {stats['compression_ratio']:.2f}x") print(f" Average compression ratio: {stats['compression_ratio']:.2f}x")
logger.info(f" Mean token length: {stats['mean_token_length']:.1f}") print(f" Mean token length: {stats['mean_token_length']:.1f}")
logger.info(f" P99 token length: {stats['p99_token_length']:.0f}") print(f" P99 token length: {stats['p99_token_length']:.0f}")
logger.info(f" Min token length: {stats['min_token_length']:.0f}") print(f" Min token length: {stats['min_token_length']:.0f}")
logger.info(f" Max token length: {stats['max_token_length']:.0f}") print(f" Max token length: {stats['max_token_length']:.0f}")
return stats return stats
@@ -386,9 +385,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
cfg: TokenizerTrainingConfig dataclass with all configuration parameters cfg: TokenizerTrainingConfig dataclass with all configuration parameters
""" """
# load dataset # load dataset
logger.info(f"Loading dataset: {cfg.repo_id}") print(f"Loading dataset: {cfg.repo_id}")
dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root) dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root)
logger.info(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames") print(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
# parse normalization mode # parse normalization mode
try: try:
@@ -398,7 +397,7 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
f"Invalid normalization_mode: {cfg.normalization_mode}. " f"Invalid normalization_mode: {cfg.normalization_mode}. "
f"Must be one of: {', '.join([m.value for m in NormalizationMode])}" f"Must be one of: {', '.join([m.value for m in NormalizationMode])}"
) from err ) from err
logger.info(f"Normalization mode: {norm_mode.value}") print(f"Normalization mode: {norm_mode.value}")
# parse encoded dimensions # parse encoded dimensions
encoded_dim_ranges = [] encoded_dim_ranges = []
@@ -407,38 +406,38 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
encoded_dim_ranges.append((start, end)) encoded_dim_ranges.append((start, end))
total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges) total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges)
logger.info(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}") print(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
# parse relative dimensions # parse relative dimensions
relative_dim_list = None relative_dim_list = None
if cfg.relative_dims is not None and cfg.relative_dims.strip(): if cfg.relative_dims is not None and cfg.relative_dims.strip():
relative_dim_list = [int(d.strip()) for d in cfg.relative_dims.split(",")] relative_dim_list = [int(d.strip()) for d in cfg.relative_dims.split(",")]
logger.info(f"Relative dimensions: {relative_dim_list}") print(f"Relative dimensions: {relative_dim_list}")
else: else:
logger.info("No relative dimensions specified") print("No relative dimensions specified")
logger.info(f"Use relative transform: {cfg.use_relative_transform}") print(f"Use relative transform: {cfg.use_relative_transform}")
if cfg.use_relative_transform and (relative_dim_list is None or len(relative_dim_list) == 0): if cfg.use_relative_transform and (relative_dim_list is None or len(relative_dim_list) == 0):
logger.warning( print(
"Warning: use_relative_transform=True but no relative_dims specified. " "Warning: use_relative_transform=True but no relative_dims specified. "
"No relative transform will be applied." "No relative transform will be applied."
) )
logger.info(f"Action horizon: {cfg.action_horizon}") print(f"Action horizon: {cfg.action_horizon}")
logger.info(f"State key: {cfg.state_key}") print(f"State key: {cfg.state_key}")
# determine episodes to process # determine episodes to process
num_episodes = dataset.num_episodes num_episodes = dataset.num_episodes
if cfg.max_episodes is not None: if cfg.max_episodes is not None:
num_episodes = min(cfg.max_episodes, num_episodes) num_episodes = min(cfg.max_episodes, num_episodes)
logger.info(f"Processing {num_episodes} episodes...") print(f"Processing {num_episodes} episodes...")
# process episodes sequentially (to avoid pickling issues with dataset) # process episodes sequentially (to avoid pickling issues with dataset)
all_chunks = [] all_chunks = []
for ep_idx in range(num_episodes): for ep_idx in range(num_episodes):
if ep_idx % 10 == 0: if ep_idx % 10 == 0:
logger.info(f" Processing episode {ep_idx}/{num_episodes}...") print(f" Processing episode {ep_idx}/{num_episodes}...")
chunks = process_episode( chunks = process_episode(
( (
@@ -456,19 +455,19 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
# concatenate all chunks # concatenate all chunks
all_chunks = np.concatenate(all_chunks, axis=0) all_chunks = np.concatenate(all_chunks, axis=0)
logger.info(f"Collected {len(all_chunks)} action chunks") print(f"Collected {len(all_chunks)} action chunks")
# extract only encoded dimensions FIRST (before normalization) # extract only encoded dimensions FIRST (before normalization)
encoded_chunks = [] encoded_chunks = []
for start, end in encoded_dim_ranges: for start, end in encoded_dim_ranges:
encoded_chunks.append(all_chunks[:, :, start:end]) encoded_chunks.append(all_chunks[:, :, start:end])
encoded_chunks = np.concatenate(encoded_chunks, axis=-1) # [N, H, D_encoded] encoded_chunks = np.concatenate(encoded_chunks, axis=-1) # [N, H, D_encoded]
logger.info(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions") print(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
# apply normalization to encoded dimensions # apply normalization to encoded dimensions
logger.info("\nBefore normalization - overall stats:") print("\nBefore normalization - overall stats:")
logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}") print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}") print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
# get normalization stats from dataset # get normalization stats from dataset
norm_stats = dataset.meta.stats norm_stats = dataset.meta.stats
@@ -490,9 +489,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
encoded_stats[stat_name] = stat_array[encoded_dim_indices] encoded_stats[stat_name] = stat_array[encoded_dim_indices]
if encoded_stats: if encoded_stats:
logger.info(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):") print(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
for stat_name, stat_values in encoded_stats.items(): for stat_name, stat_values in encoded_stats.items():
logger.info( print(
f" {stat_name}: shape={stat_values.shape}, " f" {stat_name}: shape={stat_values.shape}, "
f"range=[{np.min(stat_values):.4f}, {np.max(stat_values):.4f}]" f"range=[{np.min(stat_values):.4f}, {np.max(stat_values):.4f}]"
) )
@@ -500,27 +499,27 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
# apply normalization based on mode # apply normalization based on mode
try: try:
encoded_chunks = apply_normalization(encoded_chunks, encoded_stats, norm_mode, eps=1e-8) encoded_chunks = apply_normalization(encoded_chunks, encoded_stats, norm_mode, eps=1e-8)
logger.info(f"\nApplied {norm_mode.value} normalization") print(f"\nApplied {norm_mode.value} normalization")
except ValueError as e: except ValueError as e:
logger.warning(f"Warning: {e}. Using raw actions without normalization.") print(f"Warning: {e}. Using raw actions without normalization.")
logger.info("\nAfter normalization - overall stats:") print("\nAfter normalization - overall stats:")
logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}") print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}") print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
logger.info("\nPer-dimension stats (after normalization):") print("\nPer-dimension stats (after normalization):")
for d in range(encoded_chunks.shape[-1]): for d in range(encoded_chunks.shape[-1]):
dim_data = encoded_chunks[:, :, d] dim_data = encoded_chunks[:, :, d]
logger.info( print(
f" Dim {d}: min={np.min(dim_data):7.4f}, max={np.max(dim_data):7.4f}, " f" Dim {d}: min={np.min(dim_data):7.4f}, max={np.max(dim_data):7.4f}, "
f"mean={np.mean(dim_data):7.4f}, std={np.std(dim_data):7.4f}" f"mean={np.mean(dim_data):7.4f}, std={np.std(dim_data):7.4f}"
) )
else: else:
logger.warning("Warning: Could not extract stats for encoded dimensions, using raw actions") print("Warning: Could not extract stats for encoded dimensions, using raw actions")
else: else:
logger.warning("Warning: No normalization stats found in dataset, using raw actions") print("Warning: No normalization stats found in dataset, using raw actions")
logger.info(f"Encoded chunks shape: {encoded_chunks.shape}") print(f"Encoded chunks shape: {encoded_chunks.shape}")
# train FAST tokenizer # train FAST tokenizer
tokenizer = train_fast_tokenizer( tokenizer = train_fast_tokenizer(
@@ -562,8 +561,8 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
with open(output_path / "metadata.json", "w") as f: with open(output_path / "metadata.json", "w") as f:
json.dump(metadata, f, indent=2) json.dump(metadata, f, indent=2)
logger.info(f"\nSaved FAST tokenizer to {output_path}") print(f"\nSaved FAST tokenizer to {output_path}")
logger.info(f"Metadata: {json.dumps(metadata, indent=2)}") print(f"Metadata: {json.dumps(metadata, indent=2)}")
# push to Hugging Face Hub if requested # push to Hugging Face Hub if requested
if cfg.push_to_hub: if cfg.push_to_hub:
@@ -571,10 +570,10 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
hub_repo_id = cfg.hub_repo_id hub_repo_id = cfg.hub_repo_id
if hub_repo_id is None: if hub_repo_id is None:
hub_repo_id = output_path.name hub_repo_id = output_path.name
logger.info(f"\nNo hub_repo_id provided, using: {hub_repo_id}") print(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
logger.info(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}") print(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
logger.info(f" Private: {cfg.hub_private}") print(f" Private: {cfg.hub_private}")
try: try:
# use the tokenizer's push_to_hub method # use the tokenizer's push_to_hub method
@@ -594,15 +593,14 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
commit_message="Upload tokenizer metadata", commit_message="Upload tokenizer metadata",
) )
logger.info(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}") print(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
except Exception as e: except Exception as e:
logger.error(f"Error pushing to hub: {e}") print(f"Error pushing to hub: {e}")
logger.error(" Make sure you're logged in with `huggingface-cli login`") print(" Make sure you're logged in with `huggingface-cli login`")
def main(): def main():
"""CLI entry point that parses arguments and runs the tokenizer training.""" """CLI entry point that parses arguments and runs the tokenizer training."""
init_logging()
train_tokenizer() train_tokenizer()
@@ -171,13 +171,7 @@ class IOSPhone(BasePhone, Teleoperator):
# HEBI provides orientation in w, x, y, z format. # HEBI provides orientation in w, x, y, z format.
# Scipy's Rotation expects x, y, z, w. # Scipy's Rotation expects x, y, z, w.
quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]])) # wxyz to xyzw quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]])) # wxyz to xyzw
# ARKit can emit zero/NaN quaternions before tracking is ready or on a rot = Rotation.from_quat(quat_xyzw)
# dropped packet. Rotation.from_quat now rejects those; degrade the same
# way as a missing pose so teleop stays alive mid-session.
try:
rot = Rotation.from_quat(quat_xyzw)
except ValueError:
return False, None, None, None
pos = ar_pos - rot.apply(self.config.camera_offset) pos = ar_pos - rot.apply(self.config.camera_offset)
return True, pos, rot, pose return True, pos, rot, pose
@@ -29,12 +29,6 @@ class SOLeaderConfig:
# Whether to use degrees for angles # Whether to use degrees for angles
use_degrees: bool = True use_degrees: bool = True
# Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
# return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
# at once, which otherwise aborts the teleoperation loop. Retries are immediate (no sleep) and only
# happen on failure, so the steady-state read cost is unchanged.
num_read_retries: int = 2
@TeleoperatorConfig.register_subclass("so101_leader") @TeleoperatorConfig.register_subclass("so101_leader")
@TeleoperatorConfig.register_subclass("so100_leader") @TeleoperatorConfig.register_subclass("so100_leader")
@@ -145,7 +145,7 @@ class SOLeader(Teleoperator):
@check_if_not_connected @check_if_not_connected
def get_action(self) -> dict[str, float]: def get_action(self) -> dict[str, float]:
start = time.perf_counter() start = time.perf_counter()
action = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries) action = self.bus.sync_read("Present_Position")
action = {f"{motor}.pos": val for motor, val in action.items()} action = {f"{motor}.pos": val for motor, val in action.items()}
dt_ms = (time.perf_counter() - start) * 1e3 dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read action: {dt_ms:.1f}ms") logger.debug(f"{self} read action: {dt_ms:.1f}ms")
+7 -7
View File
@@ -41,7 +41,7 @@ class RandomSubsetApply(Transform):
def __init__( def __init__(
self, self,
transforms: Sequence[Callable[..., Any]], transforms: Sequence[Callable],
p: list[float] | None = None, p: list[float] | None = None,
n_subset: int | None = None, n_subset: int | None = None,
random_order: bool = False, random_order: bool = False,
@@ -50,7 +50,7 @@ class RandomSubsetApply(Transform):
if not isinstance(transforms, Sequence): if not isinstance(transforms, Sequence):
raise TypeError("Argument transforms should be a sequence of callables") raise TypeError("Argument transforms should be a sequence of callables")
if p is None: if p is None:
p = [1.0] * len(transforms) p = [1] * len(transforms)
elif len(p) != len(transforms): elif len(p) != len(transforms):
raise ValueError( raise ValueError(
f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}" f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}"
@@ -69,7 +69,7 @@ class RandomSubsetApply(Transform):
self.n_subset = n_subset self.n_subset = n_subset
self.random_order = random_order self.random_order = random_order
self.selected_transforms: list[Callable[..., Any]] = [] self.selected_transforms = None
def forward(self, *inputs: Any) -> Any: def forward(self, *inputs: Any) -> Any:
needs_unpacking = len(inputs) > 1 needs_unpacking = len(inputs) > 1
@@ -119,7 +119,7 @@ class SharpnessJitter(Transform):
super().__init__() super().__init__()
self.sharpness = self._check_input(sharpness) self.sharpness = self._check_input(sharpness)
def _check_input(self, sharpness: float | Sequence[float]) -> tuple[float, float]: def _check_input(self, sharpness):
if isinstance(sharpness, (int | float)): if isinstance(sharpness, (int | float)):
if sharpness < 0: if sharpness < 0:
raise ValueError("If sharpness is a single number, it must be non negative.") raise ValueError("If sharpness is a single number, it must be non negative.")
@@ -215,7 +215,7 @@ class ImageTransformsConfig:
) )
def make_transform_from_config(cfg: ImageTransformConfig) -> Transform: def make_transform_from_config(cfg: ImageTransformConfig):
if cfg.type == "SharpnessJitter": if cfg.type == "SharpnessJitter":
return SharpnessJitter(**cfg.kwargs) return SharpnessJitter(**cfg.kwargs)
@@ -236,8 +236,8 @@ class ImageTransforms(Transform):
super().__init__() super().__init__()
self._cfg = cfg self._cfg = cfg
self.weights: list[float] = [] self.weights = []
self.transforms: dict[str, Transform] = {} self.transforms = {}
for tf_name, tf_cfg in cfg.tfs.items(): for tf_name, tf_cfg in cfg.tfs.items():
if tf_cfg.weight <= 0.0: if tf_cfg.weight <= 0.0:
continue continue
+4 -13
View File
@@ -37,25 +37,16 @@ def auto_select_torch_device() -> torch.device:
# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level # 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: 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. """Given a string, return a torch.device with checks on whether the device is available."""
Raises:
ValueError: If the requested device family is known but not available on
this machine (``AssertionError`` was previously used and is easy to
mistake for a programmer bug under ``python -O`` where asserts vanish).
"""
try_device = str(try_device) try_device = str(try_device)
if try_device.startswith("cuda"): if try_device.startswith("cuda"):
if not torch.cuda.is_available(): assert torch.cuda.is_available()
raise ValueError(f"Requested device {try_device!r} but CUDA is not available.")
device = torch.device(try_device) device = torch.device(try_device)
elif try_device == "mps": elif try_device == "mps":
if not torch.backends.mps.is_available(): assert torch.backends.mps.is_available()
raise ValueError("Requested device 'mps' but MPS is not available.")
device = torch.device("mps") device = torch.device("mps")
elif try_device == "xpu": elif try_device == "xpu":
if not torch.xpu.is_available(): assert torch.xpu.is_available()
raise ValueError("Requested device 'xpu' but XPU is not available.")
device = torch.device("xpu") device = torch.device("xpu")
elif try_device == "cpu": elif try_device == "cpu":
device = torch.device("cpu") device = torch.device("cpu")
+24
View File
@@ -13,6 +13,7 @@
# limitations under the License. # limitations under the License.
import builtins import builtins
import re
from pathlib import Path from pathlib import Path
from tempfile import TemporaryDirectory from tempfile import TemporaryDirectory
from typing import Any, TypeVar from typing import Any, TypeVar
@@ -23,6 +24,29 @@ from huggingface_hub.utils import validate_hf_hub_args
from .constants import CHECKPOINTS_DIR from .constants import CHECKPOINTS_DIR
T = TypeVar("T", bound="HubMixin") T = TypeVar("T", bound="HubMixin")
REGEX_COMMIT_HASH = re.compile(r"^[0-9a-f]{40}$")
def extract_commit_hash(resolved_file: str | Path | None, revision: str | None = None) -> str | None:
"""Extract the immutable commit hash backing a resolved Hub file.
Hub cache paths contain ``snapshots/<commit_hash>/``. If the requested
revision is already a full commit hash, use it as a fallback for custom
cache layouts that do not expose the standard snapshot path.
"""
if resolved_file is not None:
path_parts = Path(resolved_file).parts
try:
snapshot_index = path_parts.index("snapshots")
commit_hash = path_parts[snapshot_index + 1]
if REGEX_COMMIT_HASH.fullmatch(commit_hash):
return commit_hash
except (ValueError, IndexError):
pass
if revision is not None and REGEX_COMMIT_HASH.fullmatch(revision):
return revision
return None
def find_latest_hub_checkpoint( def find_latest_hub_checkpoint(
+5 -5
View File
@@ -32,21 +32,21 @@ def load_json(fpath: Path) -> Any:
Returns: Returns:
Any: The data loaded from the JSON file. Any: The data loaded from the JSON file.
""" """
with open(fpath, encoding="utf-8") as f: with open(fpath) as f:
return json.load(f) return json.load(f)
def write_json(data: JsonLike, fpath: Path) -> None: def write_json(data: dict, fpath: Path) -> None:
"""Write JSON-serializable data to a file. """Write data to a JSON file.
Creates parent directories if they don't exist. Creates parent directories if they don't exist.
Args: Args:
data: JSON-serializable data to write. data (dict): The dictionary to write.
fpath (Path): The path to the output JSON file. fpath (Path): The path to the output JSON file.
""" """
fpath.parent.mkdir(exist_ok=True, parents=True) fpath.parent.mkdir(exist_ok=True, parents=True)
with open(fpath, "w", encoding="utf-8") as f: with open(fpath, "w") as f:
json.dump(data, f, indent=4, ensure_ascii=False) json.dump(data, f, indent=4, ensure_ascii=False)
-28
View File
@@ -16,39 +16,11 @@
# limitations under the License. # limitations under the License.
import logging import logging
import multiprocessing
import os import os
import signal import signal
import sys import sys
def ensure_multiprocessing_start_method(start_method: str | None) -> None:
"""Set a multiprocessing start method once, or verify the existing method matches.
Passing ``None`` leaves Python's process-wide default untouched. This is useful
when LeRobot is embedded in an application that owns multiprocessing setup.
"""
if start_method is None:
return
available_methods = multiprocessing.get_all_start_methods()
if start_method not in available_methods:
raise ValueError(
f"Multiprocessing start method must be one of {available_methods} on this platform, "
f"got {start_method!r}."
)
current_method = multiprocessing.get_start_method(allow_none=True)
if current_method is None:
multiprocessing.set_start_method(start_method)
elif current_method != start_method:
raise RuntimeError(
f"Multiprocessing start method is already {current_method!r}; cannot change it to "
f"{start_method!r}. Set the configured multiprocessing context to null to keep the "
"application's existing method, or launch LeRobot in a fresh process."
)
class ProcessSignalHandler: class ProcessSignalHandler:
"""Utility class to attach graceful shutdown signal handlers. """Utility class to attach graceful shutdown signal handlers.
-4
View File
@@ -30,10 +30,6 @@ def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: f
""" """
if seconds <= 0: if seconds <= 0:
return return
if spin_threshold < 0:
raise ValueError(f"spin_threshold must be >= 0, got {spin_threshold}")
if sleep_margin < 0:
raise ValueError(f"sleep_margin must be >= 0, got {sleep_margin}")
system = platform.system() system = platform.system()
# On macOS and Windows the scheduler / sleep granularity can make # On macOS and Windows the scheduler / sleep granularity can make
+3 -6
View File
@@ -29,13 +29,10 @@ class Rotation:
def __init__(self, quat: np.ndarray) -> None: def __init__(self, quat: np.ndarray) -> None:
"""Initialize rotation from quaternion [x, y, z, w].""" """Initialize rotation from quaternion [x, y, z, w]."""
self._quat = np.asarray(quat, dtype=float) self._quat = np.asarray(quat, dtype=float)
if self._quat.shape != (4,): # Normalize quaternion
raise ValueError(f"Quaternion must have shape (4,), got {self._quat.shape}")
# Normalize quaternion. Reject the zero vector — it has no orientation.
norm = np.linalg.norm(self._quat) norm = np.linalg.norm(self._quat)
if norm <= 0.0 or not np.isfinite(norm): if norm > 0:
raise ValueError(f"Quaternion must be a non-zero finite vector; got {self._quat} (norm={norm})") self._quat = self._quat / norm
self._quat = self._quat / norm
@classmethod @classmethod
def from_rotvec(cls, rotvec: np.ndarray) -> "Rotation": def from_rotvec(cls, rotvec: np.ndarray) -> "Rotation":
+2 -2
View File
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
from typing import NotRequired, TypedDict from typing import TypedDict
import torch import torch
@@ -28,7 +28,7 @@ class Transition(TypedDict):
next_state: dict[str, torch.Tensor] next_state: dict[str, torch.Tensor]
done: bool done: bool
truncated: bool truncated: bool
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None] complementary_info: dict[str, torch.Tensor | float | int] | None = None
def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition: def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition:
+12 -16
View File
@@ -24,6 +24,7 @@ import sys
import time import time
from collections.abc import Iterator from collections.abc import Iterator
from copy import copy, deepcopy from copy import copy, deepcopy
from datetime import datetime
from pathlib import Path from pathlib import Path
from statistics import mean from statistics import mean
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
@@ -60,16 +61,14 @@ def init_logging(
accelerator: Optional Accelerator instance (for multi-GPU detection) accelerator: Optional Accelerator instance (for multi-GPU detection)
""" """
class LeRobotFormatter(logging.Formatter): def custom_format(record: logging.LogRecord) -> str:
def format(self, record: logging.LogRecord) -> str: dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
record.lerobot_location = f"{record.pathname}:{record.lineno}"[-15:] fnameline = f"{record.pathname}:{record.lineno}"
record.lerobot_pid = f"[PID: {os.getpid()}] " if display_pid else "" pid_str = f"[PID: {os.getpid()}] " if display_pid else ""
return super().format(record) return f"{record.levelname} {pid_str}{dt} {fnameline[-15:]:>15} {record.getMessage()}"
formatter = LeRobotFormatter( formatter = logging.Formatter()
"%(levelname)s %(lerobot_pid)s%(asctime)s %(lerobot_location)15s %(message)s", formatter.format = custom_format
datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger() logger = logging.getLogger()
logger.setLevel(logging.NOTSET) logger.setLevel(logging.NOTSET)
@@ -134,13 +133,10 @@ def say(text: str, blocking: bool = False):
else: else:
raise RuntimeError("Unsupported operating system for text-to-speech.") raise RuntimeError("Unsupported operating system for text-to-speech.")
try: if blocking:
if blocking: subprocess.run(cmd, check=True)
subprocess.run(cmd, check=True, timeout=5) else:
else: subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
except (FileNotFoundError, subprocess.TimeoutExpired) as e:
logging.warning("Text-to-speech command failed: %s | Error: %s", cmd, e)
def log_say(text: str, play_sounds: bool = True, blocking: bool = False): def log_say(text: str, play_sounds: bool = True, blocking: bool = False):
+1 -68
View File
@@ -20,7 +20,7 @@
# ``` # ```
from pathlib import Path from pathlib import Path
from unittest.mock import MagicMock, patch from unittest.mock import patch
import cv2 import cv2
import numpy as np import numpy as np
@@ -123,73 +123,6 @@ def test_invalid_width_connect():
camera.connect(warmup=False) camera.connect(warmup=False)
def test_connect_cleans_up_after_settings_failure_and_allows_retry():
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
camera = OpenCVCamera(config)
opened_captures = []
def fail_settings():
opened_captures.append(camera.videocapture)
raise RuntimeError("settings failed")
with (
patch.object(camera, "_configure_capture_settings", side_effect=fail_settings),
pytest.raises(RuntimeError, match="settings failed"),
):
camera.connect(warmup=False)
assert camera.videocapture is None
assert camera.thread is None
assert not camera.is_connected
assert opened_captures[0] is not None
assert not opened_captures[0].isOpened()
camera.connect(warmup=False)
assert camera.is_connected
camera.disconnect()
def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=1)
camera = OpenCVCamera(config)
read_threads = []
def fail_warmup(*_args, **_kwargs):
read_threads.append(camera.thread)
raise TimeoutError("no frame")
with (
patch.object(camera, "async_read", side_effect=fail_warmup),
pytest.raises(TimeoutError, match="no frame"),
):
camera.connect()
assert camera.videocapture is None
assert camera.thread is None
assert not camera.is_connected
assert read_threads[0] is not None
assert not read_threads[0].is_alive()
camera.connect(warmup=False)
assert camera.is_connected
camera.disconnect()
def test_find_cameras_releases_unopened_handles():
module_path = OpenCVCamera.__module__
unopened_capture = MagicMock()
unopened_capture.isOpened.return_value = False
with (
patch(f"{module_path}.platform.system", return_value="Darwin"),
patch(f"{module_path}.MAX_OPENCV_INDEX", 1),
patch(f"{module_path}.cv2.VideoCapture", return_value=unopened_capture),
):
assert OpenCVCamera.find_cameras() == []
unopened_capture.release.assert_called_once_with()
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES) @pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
def test_read(index_or_path): def test_read(index_or_path):
config = OpenCVCameraConfig(index_or_path=index_or_path, warmup_s=0) config = OpenCVCameraConfig(index_or_path=index_or_path, warmup_s=0)
+1 -259
View File
@@ -20,7 +20,7 @@
# ``` # ```
from pathlib import Path from pathlib import Path
from unittest.mock import MagicMock, patch from unittest.mock import patch
import numpy as np import numpy as np
import pytest import pytest
@@ -30,8 +30,6 @@ from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnected
pytest.importorskip("pyrealsense2") pytest.importorskip("pyrealsense2")
import pyrealsense2 as rs
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras" TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
@@ -63,17 +61,6 @@ def test_abc_implementation():
_ = RealSenseCamera(config) _ = RealSenseCamera(config)
@pytest.mark.parametrize("option", ["exposure", "gain", "white_balance"])
def test_manual_color_option_requires_rgb(option):
with pytest.raises(ValueError, match="use_rgb=True"):
RealSenseCameraConfig(
serial_number_or_name="042",
use_rgb=False,
use_depth=True,
**{option: 100},
)
def test_connect(): def test_connect():
config = RealSenseCameraConfig(serial_number_or_name="042", warmup_s=0) config = RealSenseCameraConfig(serial_number_or_name="042", warmup_s=0)
@@ -96,27 +83,6 @@ def test_connect_invalid_camera_path(patch_realsense):
camera.connect(warmup=False) camera.connect(warmup=False)
def test_connect_cleans_up_when_sensor_configuration_fails():
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
camera = RealSenseCamera(config)
pipeline = MagicMock()
pipeline.start.return_value = MagicMock()
with (
patch("lerobot.cameras.realsense.camera_realsense.rs.pipeline", return_value=pipeline),
patch.object(camera, "_configure_rs_pipeline_config"),
patch.object(camera, "_configure_capture_settings"),
patch.object(camera, "_configure_sensor_options", side_effect=ValueError("invalid exposure")),
pytest.raises(ValueError, match="invalid exposure"),
):
camera.connect(warmup=False)
pipeline.stop.assert_called_once_with()
assert camera.rs_pipeline is None
assert camera.rs_profile is None
assert not camera.is_connected
def test_invalid_width_connect(): def test_invalid_width_connect():
config = RealSenseCameraConfig(serial_number_or_name="042", width=99999, height=480, fps=30) config = RealSenseCameraConfig(serial_number_or_name="042", width=99999, height=480, fps=30)
camera = RealSenseCamera(config) camera = RealSenseCamera(config)
@@ -125,33 +91,6 @@ def test_invalid_width_connect():
camera.connect(warmup=False) camera.connect(warmup=False)
def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30)
camera = RealSenseCamera(config)
read_threads = []
def fail_warmup(*_args, **_kwargs):
read_threads.append(camera.thread)
raise TimeoutError("no frame")
with (
patch.object(camera, "async_read", side_effect=fail_warmup),
pytest.raises(TimeoutError, match="no frame"),
):
camera.connect()
assert camera.rs_pipeline is None
assert camera.rs_profile is None
assert camera.thread is None
assert not camera.is_connected
assert read_threads[0] is not None
assert not read_threads[0].is_alive()
camera.connect(warmup=False)
assert camera.is_connected
camera.disconnect()
def test_read(): def test_read():
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0) config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
with RealSenseCamera(config) as camera: with RealSenseCamera(config) as camera:
@@ -289,203 +228,6 @@ def test_read_latest_too_old():
_ = camera.read_latest(max_age_ms=0) # immediately too old _ = camera.read_latest(max_age_ms=0) # immediately too old
def _make_mock_sensor(name: str, supported_options: set | None = None) -> MagicMock:
"""Build a fake rs.sensor that reports a name and a configurable supported-options set."""
supported = supported_options if supported_options is not None else set()
sensor = MagicMock()
sensor.get_info.return_value = name
sensor.supports.side_effect = lambda opt: opt in supported
return sensor
def _attach_mock_color_sensor(camera: RealSenseCamera, sensor: MagicMock) -> None:
"""Wire camera.rs_profile so _get_color_sensor finds the given sensor."""
profile = MagicMock()
device = MagicMock()
device.query_sensors.return_value = [sensor]
profile.get_device.return_value = device
camera.rs_profile = profile
def test_get_color_sensor_prefers_rgb_camera():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
rgb = _make_mock_sensor("RGB Camera")
stereo = _make_mock_sensor("Stereo Module")
profile = MagicMock()
device = MagicMock()
device.query_sensors.return_value = [stereo, rgb]
profile.get_device.return_value = device
camera.rs_profile = profile
assert camera._get_color_sensor() is rgb
def test_get_color_sensor_falls_back_to_stereo_module():
"""D405 has no separate RGB module; color comes from Stereo Module."""
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
stereo = _make_mock_sensor("Stereo Module")
_attach_mock_color_sensor(camera, stereo)
assert camera._get_color_sensor() is stereo
def test_get_color_sensor_raises_with_available_sensors():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
other = _make_mock_sensor("Motion Module")
_attach_mock_color_sensor(camera, other)
with pytest.raises(RuntimeError, match="Motion Module"):
camera._get_color_sensor()
def test_configure_sensor_options_skipped_when_none():
config = RealSenseCameraConfig(serial_number_or_name="042")
camera = RealSenseCamera(config)
with patch.object(RealSenseCamera, "_get_color_sensor") as mock_get:
camera._configure_sensor_options()
mock_get.assert_not_called()
def test_configure_sensor_options_applies_all_values():
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120, gain=64, white_balance=4600)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor(
"RGB Camera",
supported_options={
rs.option.enable_auto_exposure,
rs.option.exposure,
rs.option.gain,
rs.option.enable_auto_white_balance,
rs.option.white_balance,
},
)
_attach_mock_color_sensor(camera, sensor)
camera._configure_sensor_options()
sensor.set_option.assert_any_call(rs.option.enable_auto_exposure, 0)
sensor.set_option.assert_any_call(rs.option.exposure, 120)
sensor.set_option.assert_any_call(rs.option.gain, 64)
sensor.set_option.assert_any_call(rs.option.enable_auto_white_balance, 0)
sensor.set_option.assert_any_call(rs.option.white_balance, 4600)
@pytest.mark.parametrize(
("config_field", "option", "label"),
[
("exposure", rs.option.exposure, "exposure"),
("gain", rs.option.gain, "gain"),
("white_balance", rs.option.white_balance, "white balance"),
],
)
def test_configure_sensor_options_raises_when_requested_option_is_unsupported(config_field, option, label):
config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: 100})
camera = RealSenseCamera(config)
sensor = _make_mock_sensor("RGB Camera", supported_options=set())
_attach_mock_color_sensor(camera, sensor)
with pytest.raises(ValueError, match=label):
camera._configure_sensor_options()
sensor.supports.assert_any_call(option)
sensor.set_option.assert_not_called()
@pytest.mark.parametrize(
("config_field", "option", "value"),
[
("exposure", rs.option.exposure, 120),
("gain", rs.option.gain, 64),
],
)
def test_configure_sensor_options_exposure_or_gain_disables_auto_exposure(config_field, option, value):
"""white_balance=None should not touch auto white balance."""
config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: value})
camera = RealSenseCamera(config)
sensor = _make_mock_sensor(
"RGB Camera",
supported_options={rs.option.enable_auto_exposure, option},
)
_attach_mock_color_sensor(camera, sensor)
camera._configure_sensor_options()
calls = [call.args for call in sensor.set_option.call_args_list]
assert (rs.option.enable_auto_exposure, 0) in calls
assert (option, value) in calls
for opt, _ in calls:
assert opt != rs.option.enable_auto_white_balance
assert opt != rs.option.white_balance
def test_configure_sensor_options_warns_when_auto_exposure_control_is_unsupported(caplog):
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.exposure})
_attach_mock_color_sensor(camera, sensor)
with caplog.at_level("WARNING"):
camera._configure_sensor_options()
sensor.set_option.assert_called_once_with(rs.option.exposure, 120)
assert "does not support disabling auto-exposure" in caplog.text
def test_configure_sensor_options_warns_when_auto_white_balance_control_is_unsupported(caplog):
config = RealSenseCameraConfig(serial_number_or_name="042", white_balance=4600)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.white_balance})
_attach_mock_color_sensor(camera, sensor)
with caplog.at_level("WARNING"):
camera._configure_sensor_options()
sensor.set_option.assert_called_once_with(rs.option.white_balance, 4600)
assert "does not support disabling auto white balance" in caplog.text
def test_configure_sensor_options_out_of_range_raises_value_error():
"""set_option errors should be re-raised as ValueError with range diagnostics."""
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=999999)
camera = RealSenseCamera(config)
sensor = _make_mock_sensor(
"RGB Camera",
supported_options={rs.option.enable_auto_exposure, rs.option.exposure},
)
def fake_set_option(option, value):
if option == rs.option.exposure:
raise RuntimeError("value out of range")
sensor.set_option.side_effect = fake_set_option
option_range = MagicMock(min=1, max=10000, step=1, default=156)
sensor.get_option_range.return_value = option_range
_attach_mock_color_sensor(camera, sensor)
with pytest.raises(ValueError, match="exposure") as exc_info:
camera._configure_sensor_options()
msg = str(exc_info.value)
assert "999999" in msg
assert "min=1" in msg
assert "max=10000" in msg
@pytest.mark.parametrize( @pytest.mark.parametrize(
"rotation", "rotation",
[ [
@@ -0,0 +1,77 @@
# Copyright 2026 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
from dataclasses import dataclass
from lerobot.configs import PreTrainedConfig
@PreTrainedConfig.register_subclass("revision_pinning_test")
@dataclass
class RevisionPinningTestConfig(PreTrainedConfig):
@property
def observation_delta_indices(self) -> list | None:
return None
@property
def action_delta_indices(self) -> list | None:
return None
@property
def reward_delta_indices(self) -> list | None:
return None
def get_optimizer_preset(self):
raise NotImplementedError
def get_scheduler_preset(self):
raise NotImplementedError
def validate_features(self) -> None:
pass
def test_pretrained_config_pins_resolved_hub_commit(monkeypatch, tmp_path):
commit_hash = "a" * 40
snapshot_dir = tmp_path / "models--user--policy" / "snapshots" / commit_hash
RevisionPinningTestConfig(device="cpu").save_pretrained(snapshot_dir)
calls = []
def fake_hub_download(**kwargs):
calls.append(kwargs)
return str(snapshot_dir / "config.json")
monkeypatch.setattr("lerobot.configs.policies.hf_hub_download", fake_hub_download)
config = PreTrainedConfig.from_pretrained("user/policy", revision="main")
assert calls[0]["revision"] == "main"
assert config._commit_hash == commit_hash
assert config._commit_hash_source == "user/policy"
assert config.get_hub_revision("user/policy", "main") == commit_hash
assert config.get_hub_revision("user/base-policy", "base-tag") == "base-tag"
def test_runtime_commit_hash_is_not_serialized(tmp_path):
config = RevisionPinningTestConfig(device="cpu")
config._set_hub_commit_hash("a" * 40, "user/policy")
config.save_pretrained(tmp_path)
serialized_config = json.loads((tmp_path / "config.json").read_text())
assert "_commit_hash" not in serialized_config
assert "_commit_hash_source" not in serialized_config
assert "_runtime_commit_hash" not in serialized_config
assert "_runtime_commit_hash_source" not in serialized_config
@@ -1,104 +0,0 @@
# Copyright 2025 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 numpy as np
import pytest
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
from lerobot.scripts.augment_dataset_quantile_stats import (
compute_quantile_stats_for_dataset,
has_quantile_stats,
)
def _numeric_keys(dataset):
return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")]
def _image_keys(dataset):
return [k for k, v in dataset.features.items() if v["dtype"] in ("image", "video")]
def test_numeric_stats_are_unaffected_by_sampling(tmp_path, lerobot_dataset_factory):
"""Sampling only touches image/video frames; numeric features are read in
full either way, so their stats must be identical with and without sampling."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
numeric_keys = _numeric_keys(dataset)
assert numeric_keys, "fixture should expose numeric features"
for key in numeric_keys:
if key not in exact:
continue
for stat in ("mean", "std", "q01", "q50", "q99"):
if stat in exact[key]:
np.testing.assert_allclose(
sampled[key][stat],
exact[key][stat],
rtol=1e-6,
atol=1e-6,
err_msg=f"numeric feature '{key}' stat '{stat}' changed under sampling",
)
def test_image_sampling_reduces_data_but_keeps_stats_close(tmp_path, lerobot_dataset_factory):
"""For images, sampling should reduce the number of samples considered while
keeping the resulting statistics close to the exact ones."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
image_keys = _image_keys(dataset)
assert image_keys, "fixture should expose at least one image feature"
for key in image_keys:
# sampling actually looked at fewer pixels
assert sampled[key]["count"][0] < exact[key]["count"][0]
# but per-channel mean stays close
np.testing.assert_allclose(
sampled[key]["mean"],
exact[key]["mean"],
rtol=0.15,
err_msg=f"image feature '{key}' mean drifted too far under sampling",
)
def test_short_episodes_use_all_frames(tmp_path, lerobot_dataset_factory):
"""With episodes shorter than the sampling floor, sampling is a no-op and
must produce exactly the same stats as the exact path."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=40, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
for key in _image_keys(dataset):
assert sampled[key]["count"][0] == exact[key]["count"][0]
def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory):
"""The computed stats should contain quantile keys for the dataset."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=200, use_videos=False
)
stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
assert has_quantile_stats(stats)
-14
View File
@@ -482,20 +482,6 @@ def test_add_frame_works_in_write_mode(tmp_path):
# ── Resume mode ────────────────────────────────────────────────────── # ── Resume mode ──────────────────────────────────────────────────────
def test_resume_freshly_created_empty_dataset(tmp_path):
"""resume() accepts a local dataset created before any episode was recorded."""
root = tmp_path / "resume_empty_ds"
LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root)
resumed = LeRobotDataset.resume(repo_id=DUMMY_REPO_ID, root=root)
assert isinstance(resumed.writer, DatasetWriter)
assert resumed.meta.total_episodes == 0
assert resumed.meta.total_frames == 0
assert resumed.meta.tasks is None
assert resumed.meta.episodes is None
def test_resume_creates_writer(tmp_path): def test_resume_creates_writer(tmp_path):
"""After resume(), writer is a DatasetWriter.""" """After resume(), writer is a DatasetWriter."""
root = tmp_path / "resume_ds" root = tmp_path / "resume_ds"
-13
View File
@@ -294,19 +294,6 @@ def test__sync_read(addr, length, ids_values, mock_motors, dummy_motors):
assert read_values == ids_values assert read_values == ids_values
def test__sync_read_retries_after_transient_failure(mock_motors, dummy_motors):
addr, length, ids_values = (10, 4, {1: 1337})
stub = mock_motors.build_sync_read_stub(addr, length, ids_values, num_invalid_try=1)
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
bus.connect(handshake=False)
read_values, read_comm = bus._sync_read(addr, length, list(ids_values), num_retry=1)
assert read_comm == scs.COMM_SUCCESS
assert read_values == ids_values
assert mock_motors.stubs[stub].calls == 2
@pytest.mark.parametrize("raise_on_error", (True, False)) @pytest.mark.parametrize("raise_on_error", (True, False))
def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors): def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors):
addr, length, ids_values = (10, 4, {1: 1337}) addr, length, ids_values = (10, 4, {1: 1337})
-54
View File
@@ -496,60 +496,6 @@ def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
assert "embodiment_id" in processed assert "embodiment_id" in processed
def test_reconcile_evo1_processors_repads_overridden_stats(tmp_path):
"""Loading a checkpoint and injecting raw (unpadded) dataset stats must be re-padded.
Regression test: lerobot-train passes the raw dataset stats as normalizer/unnormalizer
overrides when resuming from a checkpoint (e.g. stage2 from a stage1 checkpoint). Those stats
are at the dataset dims (e.g. LIBERO state=8/action=7), but EVO1 pads state/action to
max_state_dim/max_action_dim before normalization, so reconcile_evo1_processors must re-pad the
stats or normalization crashes with a shape mismatch.
"""
config = make_config()
preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=make_stats())
preprocessor.save_pretrained(tmp_path)
postprocessor.save_pretrained(tmp_path)
# Reload with the generic override path injecting raw, unpadded dataset stats.
raw_stats = make_stats()
loaded_pre = PolicyProcessorPipeline.from_pretrained(
tmp_path,
config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
overrides={"normalizer_processor": {"stats": raw_stats}},
to_transition=batch_to_transition,
to_output=transition_to_batch,
)
loaded_post = PolicyProcessorPipeline.from_pretrained(
tmp_path,
config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
overrides={"unnormalizer_processor": {"stats": raw_stats}},
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
)
# Sanity: the override really injected unpadded stats before reconciliation.
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (STATE_DIM,)
loaded_pre, loaded_post = reconcile_evo1_processors(config, loaded_pre, loaded_post)
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
unnormalizer = next(step for step in loaded_post.steps if isinstance(step, UnnormalizerProcessorStep))
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
assert normalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
assert unnormalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
# Normalizing a padded state must not raise (this is the exact runtime path that crashed).
processed = loaded_pre(
{
"task": "pick the block",
OBS_STATE: torch.zeros(STATE_DIM),
f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
}
)
assert processed[OBS_STATE].shape == (1, MAX_STATE_DIM)
def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch): def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model) monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
policy = modeling_evo1.Evo1Policy(make_config()) policy = modeling_evo1.Evo1Policy(make_config())
@@ -1,83 +0,0 @@
# Copyright 2026 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 types import SimpleNamespace
from unittest.mock import MagicMock
import torch
import lerobot.policies.factory as policy_factory
def test_make_policy_keeps_peft_adapter_and_base_revisions_separate(monkeypatch):
cfg = SimpleNamespace(
type="mock",
device="cpu",
pretrained_path="user/adapter",
pretrained_revision="adapter-sha",
use_peft=True,
input_features={},
output_features={},
)
dataset_meta = SimpleNamespace(features={}, stats={})
base_policy = torch.nn.Linear(1, 1)
policy_from_pretrained = MagicMock(return_value=base_policy)
policy_class = SimpleNamespace(from_pretrained=policy_from_pretrained)
monkeypatch.setattr(policy_factory, "get_policy_class", lambda _: policy_class)
monkeypatch.setattr(policy_factory, "dataset_to_policy_features", lambda _: {})
monkeypatch.setattr(policy_factory, "validate_visual_features_consistency", lambda *args: None)
peft_config = SimpleNamespace(
base_model_name_or_path="user/base-policy",
revision="base-sha",
)
peft_config_from_pretrained = MagicMock(return_value=peft_config)
adapted_policy = torch.nn.Linear(1, 1)
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
require_package = MagicMock()
monkeypatch.setattr(policy_factory, "require_package", require_package)
monkeypatch.setattr(
policy_factory,
"PeftConfig",
SimpleNamespace(from_pretrained=peft_config_from_pretrained),
)
monkeypatch.setattr(
policy_factory,
"PeftModel",
SimpleNamespace(from_pretrained=peft_model_from_pretrained),
)
policy = policy_factory.make_policy(cfg, ds_meta=dataset_meta)
assert policy is adapted_policy
require_package.assert_called_once_with("peft", extra="peft")
peft_config_from_pretrained.assert_called_once_with(
"user/adapter",
revision="adapter-sha",
)
policy_from_pretrained.assert_called_once_with(
config=cfg,
dataset_stats=dataset_meta.stats,
dataset_meta=dataset_meta,
pretrained_name_or_path="user/base-policy",
revision="base-sha",
)
peft_model_from_pretrained.assert_called_once_with(
base_policy,
"user/adapter",
config=peft_config,
revision="adapter-sha",
is_trainable=True,
)
@@ -113,7 +113,6 @@ def test_gaussian_actor_config_default_initialization():
# Concurrency configuration # Concurrency configuration
assert config.concurrency.actor == "threads" assert config.concurrency.actor == "threads"
assert config.concurrency.learner == "threads" assert config.concurrency.learner == "threads"
assert config.concurrency.multiprocessing_context == "spawn"
assert isinstance(config.actor_network_kwargs, ActorNetworkConfig) assert isinstance(config.actor_network_kwargs, ActorNetworkConfig)
assert isinstance(config.policy_kwargs, PolicyConfig) assert isinstance(config.policy_kwargs, PolicyConfig)
@@ -153,7 +152,6 @@ def test_concurrency_config():
config = ConcurrencyConfig() config = ConcurrencyConfig()
assert config.actor == "threads" assert config.actor == "threads"
assert config.learner == "threads" assert config.learner == "threads"
assert config.multiprocessing_context == "spawn"
def test_gaussian_actor_config_custom_initialization(): def test_gaussian_actor_config_custom_initialization():
@@ -0,0 +1,122 @@
# Copyright 2026 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.policies import factory
from lerobot.policies.act.configuration_act import ACTConfig
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.processor import PolicyProcessorPipeline
class _MinimalPolicy(PreTrainedPolicy):
config_class = ACTConfig
name = "minimal_revision_test"
def get_optim_params(self) -> dict:
return {}
def reset(self) -> None:
pass
def forward(self, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, dict | None]:
return torch.tensor(0), None
def predict_action_chunk(self, batch: dict[str, torch.Tensor], **kwargs) -> torch.Tensor:
return torch.tensor(0)
def select_action(self, batch: dict[str, torch.Tensor], **kwargs) -> torch.Tensor:
return torch.tensor(0)
def _skip_safetensor_loading(monkeypatch):
monkeypatch.setattr(
_MinimalPolicy,
"_load_as_safetensor",
classmethod(lambda cls, model, model_file, map_location, strict: model),
)
def test_policy_weights_use_config_commit_hash(monkeypatch):
config = ACTConfig(device="cpu")
config._set_hub_commit_hash("a" * 40, "user/policy")
calls = []
def fake_hub_download(**kwargs):
calls.append(kwargs)
return "/unused/model.safetensors"
monkeypatch.setattr("lerobot.policies.pretrained.hf_hub_download", fake_hub_download)
_skip_safetensor_loading(monkeypatch)
_MinimalPolicy.from_pretrained("user/policy", config=config, revision="main")
assert calls[0]["revision"] == "a" * 40
def test_policy_does_not_reuse_commit_hash_for_another_repo(monkeypatch):
config = ACTConfig(device="cpu")
config._set_hub_commit_hash("a" * 40, "user/adapter")
calls = []
def fake_hub_download(**kwargs):
calls.append(kwargs)
return "/unused/model.safetensors"
monkeypatch.setattr("lerobot.policies.pretrained.hf_hub_download", fake_hub_download)
_skip_safetensor_loading(monkeypatch)
_MinimalPolicy.from_pretrained("user/base-policy", config=config, revision="base-tag")
assert calls[0]["revision"] == "base-tag"
def test_policy_records_weight_commit_for_explicit_config(monkeypatch):
commit_hash = "a" * 40
config = ACTConfig(device="cpu")
def fake_hub_download(**kwargs):
return f"/cache/models--user--policy/snapshots/{commit_hash}/model.safetensors"
monkeypatch.setattr("lerobot.policies.pretrained.hf_hub_download", fake_hub_download)
_skip_safetensor_loading(monkeypatch)
_MinimalPolicy.from_pretrained("user/policy", config=config, revision="main")
assert config._commit_hash == commit_hash
assert config._commit_hash_source == "user/policy"
def test_processor_factory_uses_config_commit_hash(monkeypatch):
config = ACTConfig(device="cpu")
config._set_hub_commit_hash("a" * 40, "user/policy")
calls = []
def fake_from_pretrained(cls, **kwargs):
calls.append(kwargs)
return PolicyProcessorPipeline(steps=[])
monkeypatch.setattr(
factory.PolicyProcessorPipeline,
"from_pretrained",
classmethod(fake_from_pretrained),
)
factory.make_pre_post_processors(
config,
pretrained_path="user/policy",
pretrained_revision="main",
)
assert [call["revision"] for call in calls] == ["a" * 40, "a" * 40]
@@ -241,6 +241,68 @@ def test_from_pretrained_hub_source_missing_local_state_still_calls_hub(monkeypa
ProcessorStepRegistry.unregister("hub_state_step") ProcessorStepRegistry.unregister("hub_state_step")
def test_from_pretrained_pins_hub_state_to_config_commit(monkeypatch, tmp_path):
"""A mutable processor revision is resolved once and reused for state files."""
@ProcessorStepRegistry.register("pinned_hub_state_step")
class PinnedHubStateStep(ProcessorStep):
def __init__(self):
self.value = torch.tensor(0)
def __call__(self, transition: EnvTransition) -> EnvTransition:
return transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
self.value = state["value"]
try:
commit_hash = "a" * 40
snapshot_dir = tmp_path / "models--user--policy" / "snapshots" / commit_hash
snapshot_dir.mkdir(parents=True)
config_path = snapshot_dir / "processor.json"
config_path.write_text(
json.dumps(
{
"name": "PinnedHubStatePipeline",
"steps": [
{
"registry_name": "pinned_hub_state_step",
"state_file": "hub_state.safetensors",
}
],
}
)
)
state_path = tmp_path / "downloaded.safetensors"
save_file({"value": torch.tensor(7)}, state_path)
calls = []
def fake_hub_download(**kwargs):
calls.append(kwargs)
if kwargs["filename"] == "processor.json":
return str(config_path)
return str(state_path)
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fake_hub_download)
pipeline = DataProcessorPipeline.from_pretrained(
"user/policy",
config_filename="processor.json",
revision="main",
)
assert calls[0]["revision"] == "main"
assert calls[1]["revision"] == commit_hash
assert pipeline.steps[0].value.item() == 7
finally:
ProcessorStepRegistry.unregister("pinned_hub_state_step")
# Config Validation Tests # Config Validation Tests
@@ -1,45 +0,0 @@
#!/usr/bin/env python
# Copyright 2025 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 pytest
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.robots.so_follower.robot_kinematic_processor import (
ForwardKinematicsJointsToEEAction,
ForwardKinematicsJointsToEEObservation,
)
MOTOR_NAMES = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
EE_KEYS = {f"ee.{k}" for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]}
def _joint_bucket(feature_type: FeatureType) -> dict[str, PolicyFeature]:
return {f"{n}.pos": PolicyFeature(type=feature_type, shape=(1,)) for n in MOTOR_NAMES}
@pytest.mark.parametrize(
("step_cls", "bucket", "feature_type"),
[
(ForwardKinematicsJointsToEEAction, PipelineFeatureType.ACTION, FeatureType.ACTION),
(ForwardKinematicsJointsToEEObservation, PipelineFeatureType.OBSERVATION, FeatureType.STATE),
],
)
def test_fk_feature_schema(step_cls, bucket, feature_type):
features = {PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: {}}
features[bucket] = _joint_bucket(feature_type)
out = step_cls(kinematics=None, motor_names=MOTOR_NAMES).transform_features(features)[bucket]
assert set(out) == EE_KEYS
assert {feature.type for feature in out.values()} == {feature_type}
+2 -23
View File
@@ -49,7 +49,7 @@ def _make_bus_mock() -> MagicMock:
@pytest.fixture @pytest.fixture
def follower(tmp_path): def follower():
bus_mock = _make_bus_mock() bus_mock = _make_bus_mock()
def _bus_side_effect(*_args, **kwargs): def _bus_side_effect(*_args, **kwargs):
@@ -71,7 +71,7 @@ def follower(tmp_path):
), ),
patch.object(SO100Follower, "configure", lambda self: None), patch.object(SO100Follower, "configure", lambda self: None),
): ):
cfg = SO100FollowerConfig(port="/dev/null", calibration_dir=tmp_path) cfg = SO100FollowerConfig(port="/dev/null")
robot = SO100Follower(cfg) robot = SO100Follower(cfg)
yield robot yield robot
if robot.is_connected: if robot.is_connected:
@@ -99,27 +99,6 @@ def test_get_observation(follower):
assert obs[f"{motor}.pos"] == idx assert obs[f"{motor}.pos"] == idx
def test_get_observation_uses_read_retries(follower):
# Feetech buses can intermittently fail a sync_read; the follower should forward the configured
# retry count so transient failures don't abort the control loop (see #3131).
follower.config.num_read_retries = 7
follower.connect()
follower.get_observation()
follower.bus.sync_read.assert_called_once_with("Present_Position", num_retry=7)
def test_send_action_uses_read_retries(follower):
follower.config.max_relative_target = 10.0
follower.config.num_read_retries = 7
follower.connect()
action = {f"{motor}.pos": value * 10 for value, motor in enumerate(follower.bus.motors, 1)}
follower.send_action(action)
follower.bus.sync_read.assert_called_once_with("Present_Position", num_retry=7)
def test_send_action(follower): def test_send_action(follower):
follower.connect() follower.connect()
+7 -14
View File
@@ -185,25 +185,18 @@ def test_load_pretrained_peft_policy_keeps_adapter_and_base_revisions_separate(m
peft_config_from_pretrained = MagicMock(return_value=peft_config) peft_config_from_pretrained = MagicMock(return_value=peft_config)
adapted_policy = MagicMock() adapted_policy = MagicMock()
peft_model_from_pretrained = MagicMock(return_value=adapted_policy) peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
require_package = MagicMock() monkeypatch.setitem(
monkeypatch.setattr(rollout_context, "require_package", require_package) sys.modules,
monkeypatch.setattr( "peft",
rollout_context, SimpleNamespace(
"PeftConfig", PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained),
SimpleNamespace(from_pretrained=peft_config_from_pretrained), PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained),
raising=False, ),
)
monkeypatch.setattr(
rollout_context,
"PeftModel",
SimpleNamespace(from_pretrained=peft_model_from_pretrained),
raising=False,
) )
policy = rollout_context._load_pretrained_policy(policy_config) policy = rollout_context._load_pretrained_policy(policy_config)
assert policy is adapted_policy assert policy is adapted_policy
require_package.assert_called_once_with("peft", extra="peft")
peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha") peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha")
policy_class.from_pretrained.assert_called_once_with( policy_class.from_pretrained.assert_called_once_with(
pretrained_name_or_path="user/base-policy", pretrained_name_or_path="user/base-policy",
-36
View File
@@ -1,36 +0,0 @@
#!/usr/bin/env python
# Copyright 2026 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 unittest.mock import patch
import pytest
import torch
from lerobot.utils.device_utils import get_safe_torch_device, is_torch_device_available
def test_cpu_always_available():
assert get_safe_torch_device("cpu") == torch.device("cpu")
assert is_torch_device_available("cpu")
def test_missing_cuda_raises_valueerror():
with patch("torch.cuda.is_available", return_value=False), pytest.raises(ValueError, match="CUDA"):
get_safe_torch_device("cuda")
def test_missing_mps_raises_valueerror():
with patch("torch.backends.mps.is_available", return_value=False), pytest.raises(ValueError, match="MPS"):
get_safe_torch_device("mps")
+18 -1
View File
@@ -14,7 +14,7 @@
from unittest.mock import MagicMock from unittest.mock import MagicMock
from lerobot.utils.hub import find_latest_hub_checkpoint from lerobot.utils.hub import extract_commit_hash, find_latest_hub_checkpoint
def _patch_list_files(monkeypatch, files): def _patch_list_files(monkeypatch, files):
@@ -52,3 +52,20 @@ def test_find_latest_hub_checkpoint_ignores_non_step_entries(monkeypatch):
def test_find_latest_hub_checkpoint_none_when_no_checkpoints(monkeypatch): def test_find_latest_hub_checkpoint_none_when_no_checkpoints(monkeypatch):
_patch_list_files(monkeypatch, ["config.json", "model.safetensors"]) _patch_list_files(monkeypatch, ["config.json", "model.safetensors"])
assert find_latest_hub_checkpoint("u/run") is None assert find_latest_hub_checkpoint("u/run") is None
def test_extract_commit_hash_from_hub_snapshot_path():
commit_hash = "a" * 40
resolved_file = f"/cache/models--user--policy/snapshots/{commit_hash}/config.json"
assert extract_commit_hash(resolved_file, revision="main") == commit_hash
def test_extract_commit_hash_falls_back_to_full_sha_revision():
commit_hash = "b" * 40
assert extract_commit_hash("/custom/cache/config.json", revision=commit_hash) == commit_hash
def test_extract_commit_hash_rejects_mutable_revision_without_snapshot():
assert extract_commit_hash("/custom/cache/config.json", revision="main") is None
-46
View File
@@ -1,46 +0,0 @@
#!/usr/bin/env python
# Copyright 2026 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 numpy as np
import pytest
from lerobot.utils.rotation import Rotation
def test_zero_quaternion_rejected():
with pytest.raises(ValueError, match="non-zero"):
Rotation(np.zeros(4))
def test_non_finite_quaternion_rejected():
with pytest.raises(ValueError, match="non-zero|finite"):
Rotation(np.array([np.nan, 0.0, 0.0, 1.0]))
def test_wrong_shape_rejected():
with pytest.raises(ValueError, match="shape"):
Rotation(np.array([1.0, 0.0, 0.0]))
def test_identity_roundtrip():
r = Rotation.from_rotvec(np.zeros(3))
assert np.allclose(r.as_rotvec(), 0.0)
assert np.allclose(r.as_matrix(), np.eye(3))
def test_rotvec_roundtrip():
rotvec = np.array([0.1, -0.2, 0.3])
r = Rotation.from_rotvec(rotvec)
assert np.allclose(r.as_rotvec(), rotvec, atol=1e-6)
-230
View File
@@ -1,230 +0,0 @@
#!/usr/bin/env python3
"""Provision the SONIC decoder checkpoint at ``lerobot/sonic_decoder``.
Takes NVIDIA's ``nvidia/GEAR-SONIC/model_decoder.onnx``, embeds the SONIC deploy constants
(``kp``/``kd`` PD gains, ``default_angles`` standing pose, the residual ``action_scale``, and
the ``neutral_token`` idle latent) into the ONNX ``metadata_props`` (the convention Holosoma
uses for its gains), and pushes the result to ``lerobot/sonic_decoder``. After this runs, the
runtime loads the decoder *and* every one of these constants straight from the checkpoint --
no motor-physics math at deploy time, so ``sonic_whole_body.py`` carries none of the
armature/bandwidth machinery nor any hardcoded deploy constants.
The constants here are derived once from Unitree motor physics (armature + target bandwidth).
That derivation is intentionally kept in this one-off provisioning script (not the runtime);
the shared/harmonic helper is a separate PR.
Build only (no network/auth needed if the source ONNX is already cached):
python upload_sonic_decoder.py --out ./sonic_decoder
Build + upload:
huggingface-cli login # or export HF_TOKEN=...
python upload_sonic_decoder.py --upload
"""
from __future__ import annotations
import argparse
import json
import pathlib
import numpy as np
import onnx
from huggingface_hub import hf_hub_download
SRC_REPO_ID = "nvidia/GEAR-SONIC"
SRC_FILENAME = "model_decoder.onnx"
DST_REPO_ID = "lerobot/sonic_decoder"
# ── SONIC deploy-constant derivation (provisioning-time only) ─────────────────
# All constants are (29,) in IsaacLab joint order: legs, waist, arms.
# kp = armature * w**2, kd = 4 * armature * w, with a x2 factor on the stiff joints
# (ankles + waist). action_scale = 0.25 * effort / (armature * w**2) is the residual
# scaling that maps decoder output to a joint-angle delta on top of default_angles.
NATURAL_FREQ = 10.0 * 2.0 * np.pi
MOTOR_ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425}
EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0}
MOTOR_MODELS = (
["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
+ ["7520_14", "5020", "5020"]
+ ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
)
DOUBLE_INDICES = {4, 5, 10, 11, 13, 14} # ankles + waist
# Nominal standing pose (rad), 29 joints in IsaacLab order. Decoder actions are residuals
# added on top of this.
DEFAULT_ANGLES = [
-0.312,
0.0,
0.0,
0.669,
-0.363,
0.0, # left leg
-0.312,
0.0,
0.0,
0.669,
-0.363,
0.0, # right leg
0.0,
0.0,
0.0, # waist
0.2,
0.2,
0.0,
0.6,
0.0,
0.0,
0.0, # left arm
0.2,
-0.2,
0.0,
0.6,
0.0,
0.0,
0.0, # right arm
]
# Neutral idle token (64-D), held until the first real token arrives. Captured from the
# encoder while the robot stood idle in sim: the encoder is an FSQ bottleneck (~5 bit/dim,
# Div(16)), so tokens live on the 1/16 grid. We store the integer FSQ codes and rescale by
# 1/16 -> an exact on-grid token that decodes to a stable, natural standing pose (unlike the
# literal all-zero token, which is off-manifold and decodes to a slightly goofy stance).
NEUTRAL_TOKEN_CODES = [
-1,
3,
1,
-1,
1,
-3,
6,
1,
1,
1,
-2,
-4,
-2,
0,
-3,
-1,
2,
-1,
-3,
-5,
3,
1,
1,
-4,
-1,
-1,
1,
-7,
0,
1,
2,
-2,
5,
-2,
-2,
-4,
0,
-1,
3,
-1,
0,
-5,
-1,
0,
-4,
0,
0,
-1,
-1,
2,
-2,
1,
3,
3,
1,
0,
0,
6,
0,
-7,
3,
0,
2,
-2,
]
def compute_kp_kd() -> tuple[list[float], list[float]]:
"""Return (kp, kd) as plain float lists, (29,) in IsaacLab joint order."""
def stiffness(k):
return MOTOR_ARMATURE[k] * NATURAL_FREQ**2
def damping(k):
return 4.0 * MOTOR_ARMATURE[k] * NATURAL_FREQ
kp = [(2 if i in DOUBLE_INDICES else 1) * stiffness(k) for i, k in enumerate(MOTOR_MODELS)]
kd = [(2 if i in DOUBLE_INDICES else 1) * damping(k) for i, k in enumerate(MOTOR_MODELS)]
return kp, kd
def compute_action_scale() -> list[float]:
"""Return the per-joint residual action scale, (29,) in IsaacLab joint order."""
return [0.25 * EFFORT[k] / (MOTOR_ARMATURE[k] * NATURAL_FREQ**2) for k in MOTOR_MODELS]
def build(out_dir: pathlib.Path) -> pathlib.Path:
"""Download the source decoder, embed the deploy-constant metadata, save to ``out_dir``."""
src = hf_hub_download(repo_id=SRC_REPO_ID, filename=SRC_FILENAME)
model = onnx.load(src)
kp, kd = compute_kp_kd()
neutral_token = [c / 16.0 for c in NEUTRAL_TOKEN_CODES] # FSQ Div(16): codes -> on-grid token
meta = {prop.key: prop.value for prop in model.metadata_props}
meta["kp"] = json.dumps(kp)
meta["kd"] = json.dumps(kd)
meta["action_scale"] = json.dumps(compute_action_scale())
meta["default_angles"] = json.dumps(DEFAULT_ANGLES)
meta["neutral_token"] = json.dumps(neutral_token)
# Rewrite metadata_props with the merged dict.
del model.metadata_props[:]
for key, value in meta.items():
model.metadata_props.add(key=key, value=value)
out_dir.mkdir(parents=True, exist_ok=True)
out_path = out_dir / SRC_FILENAME
onnx.save(model, out_path)
print(f"Wrote {out_path} with kp/kd/action_scale/default_angles/neutral_token metadata.")
return out_path
def upload(out_path: pathlib.Path) -> None:
from huggingface_hub import HfApi
api = HfApi()
api.create_repo(repo_id=DST_REPO_ID, repo_type="model", exist_ok=True)
api.upload_file(
path_or_fileobj=str(out_path),
path_in_repo=SRC_FILENAME,
repo_id=DST_REPO_ID,
repo_type="model",
)
print(f"Uploaded {out_path.name} -> {DST_REPO_ID}")
def main() -> None:
p = argparse.ArgumentParser()
p.add_argument("--out", type=pathlib.Path, default=pathlib.Path("./sonic_decoder"))
p.add_argument("--upload", action="store_true", help="Push the built ONNX to the hub")
args = p.parse_args()
out_path = build(args.out)
if args.upload:
upload(out_path)
if __name__ == "__main__":
main()