mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 1834f819a6 | |||
| 741005d719 |
@@ -33,7 +33,7 @@ jobs:
|
||||
github.event.workflow_run.event == 'pull_request' &&
|
||||
github.event.workflow_run.conclusion == 'success' &&
|
||||
github.repository == 'huggingface/lerobot'
|
||||
uses: huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml@931031bf2b54aabb134ceb54980a6a2860a00f11 # main
|
||||
uses: huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml@6108e850ae1cf2f71bb0815a600bcd50c39abfa7 # main
|
||||
with:
|
||||
package_name: lerobot
|
||||
secrets:
|
||||
|
||||
@@ -60,14 +60,10 @@ jobs:
|
||||
github.repository == 'huggingface/lerobot'
|
||||
permissions:
|
||||
contents: read
|
||||
uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@931031bf2b54aabb134ceb54980a6a2860a00f11 # main
|
||||
uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@6108e850ae1cf2f71bb0815a600bcd50c39abfa7 # main
|
||||
with:
|
||||
commit_sha: ${{ github.sha }}
|
||||
package: lerobot
|
||||
# The shared workflow builds its venv with the runner's system Python, which is 3.10 on
|
||||
# ubuntu-22.04. lerobot requires >=3.12, so without this the install fails during setup —
|
||||
# before `pre_command` below ever runs. Added upstream in huggingface/doc-builder#808.
|
||||
python_version: "3.12"
|
||||
# doc-builder ships a mock-deps registry entry for lerobot, so the reusable workflow takes its
|
||||
# "light install" path: `pip install ./lerobot --no-deps` plus a handful of real dependencies.
|
||||
# That is not enough to import lerobot — draccus runs `register_subclass` at import time and
|
||||
@@ -96,12 +92,11 @@ jobs:
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@931031bf2b54aabb134ceb54980a6a2860a00f11 # main
|
||||
uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@6108e850ae1cf2f71bb0815a600bcd50c39abfa7 # main
|
||||
with:
|
||||
commit_sha: ${{ github.event.pull_request.head.sha }}
|
||||
pr_number: ${{ github.event.number }}
|
||||
package: lerobot
|
||||
# See the comment on build_main_docs. The PR workflow passes its own `--version pr_<n>`, so no
|
||||
# additional_args are needed here.
|
||||
python_version: "3.12"
|
||||
pre_command: uv pip install "./lerobot[dataset]"
|
||||
|
||||
@@ -215,4 +215,6 @@
|
||||
title: Environments
|
||||
- local: api/configs
|
||||
title: Configuration
|
||||
- local: api/rl
|
||||
title: Reinforcement Learning
|
||||
title: "API Reference"
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
# Reinforcement Learning
|
||||
|
||||
`lerobot.rl` is the distributed actor/learner reinforcement-learning stack behind
|
||||
[Train a Robot with RL](../hilserl) (HIL-SERL) and [Train RL in Simulation](../hilserl_sim). Algorithms,
|
||||
the replay buffer, data sources, and the trainer are gRPC-free and usable standalone; the actor/learner
|
||||
entry points (`actor`, `learner`, `learner_service`) additionally require `pip install 'lerobot[hilserl]'`.
|
||||
|
||||
## TrainRLServerPipelineConfig
|
||||
|
||||
Top-level configuration for both the `lerobot-actor` and `lerobot-learner` CLIs.
|
||||
|
||||
[[autodoc]] lerobot.rl.train_rl.TrainRLServerPipelineConfig
|
||||
|
||||
## RLAlgorithm
|
||||
|
||||
Abstract base every RL algorithm subclasses.
|
||||
|
||||
[[autodoc]] lerobot.rl.algorithms.base.RLAlgorithm
|
||||
|
||||
## RLAlgorithmConfig
|
||||
|
||||
[[autodoc]] lerobot.rl.algorithms.configs.RLAlgorithmConfig
|
||||
|
||||
## TrainingStats
|
||||
|
||||
[[autodoc]] lerobot.rl.algorithms.configs.TrainingStats
|
||||
|
||||
## make_algorithm
|
||||
|
||||
[[autodoc]] lerobot.rl.algorithms.factory.make_algorithm
|
||||
|
||||
## make_algorithm_config
|
||||
|
||||
[[autodoc]] lerobot.rl.algorithms.factory.make_algorithm_config
|
||||
|
||||
## get_algorithm_class
|
||||
|
||||
[[autodoc]] lerobot.rl.algorithms.factory.get_algorithm_class
|
||||
|
||||
## SAC
|
||||
|
||||
[[autodoc]] lerobot.rl.algorithms.sac.sac_algorithm.SACAlgorithm
|
||||
- all
|
||||
|
||||
[[autodoc]] lerobot.rl.algorithms.sac.configuration_sac.SACAlgorithmConfig
|
||||
|
||||
## ReplayBuffer
|
||||
|
||||
In-memory replay buffer of transitions, sampled in batches for off-policy training.
|
||||
|
||||
[[autodoc]] lerobot.rl.buffer.ReplayBuffer
|
||||
- all
|
||||
|
||||
[[autodoc]] lerobot.rl.buffer.BatchTransition
|
||||
|
||||
## DataMixer
|
||||
|
||||
Abstract interface for combining online and offline data sources into training batches.
|
||||
|
||||
[[autodoc]] lerobot.rl.data_sources.DataMixer
|
||||
- all
|
||||
|
||||
[[autodoc]] lerobot.rl.data_sources.OnlineOfflineMixer
|
||||
- all
|
||||
|
||||
## RLTrainer
|
||||
|
||||
Unified training-step orchestrator: holds the algorithm, a `DataMixer`, and an optional preprocessor.
|
||||
|
||||
[[autodoc]] lerobot.rl.trainer.RLTrainer
|
||||
- all
|
||||
|
||||
## Actor / learner CLIs
|
||||
|
||||
The distributed actor and learner processes communicate over gRPC; see [Train a Robot with
|
||||
RL](../hilserl) for the full workflow.
|
||||
|
||||
[[autodoc]] lerobot.rl.actor.actor_cli
|
||||
|
||||
[[autodoc]] lerobot.rl.learner.train_cli
|
||||
|
||||
[[autodoc]] lerobot.rl.learner_service.LearnerService
|
||||
- all
|
||||
|
||||
## eval_policy
|
||||
|
||||
[[autodoc]] lerobot.rl.eval_policy.eval_policy
|
||||
@@ -62,10 +62,7 @@ Reference data points on a 4×H100 80 GB cluster (`accelerate launch --num_proce
|
||||
| `smolvla` | 27m 49s | 0.312 | 0.011 | ~80% | `--policy.path=lerobot/smolvla_base`, `freeze_vision_encoder=false`, `train_expert_only=false` |
|
||||
| `pi05` | 3h 41m | 2.548 | 0.014 | ~95% | `--policy.pretrained_path=lerobot/pi05_base`, `gradient_checkpointing=true`, `dtype=bfloat16`, vision encoder + expert trained |
|
||||
|
||||
Training logs separate the full iteration into `dataloading_s` (`next(dl_iter)`), `preprocessing_s`
|
||||
(image conversion and the policy pipeline), and `update_s` (the optimizer update). `step_s` covers all
|
||||
three and drives `samples_per_s`. The benchmark above predates this split, so its `dataloading_s` includes
|
||||
preprocessing.
|
||||
The `dataloading_s` vs. `update_s` ratio is the diagnostic that matters: when `dataloading_s` approaches `update_s`, more GPUs stop helping — your dataloader is the bottleneck and you should look at `--num_workers`, image resolution, and disk speed before adding compute.
|
||||
|
||||
### Schedule and checkpoints
|
||||
|
||||
|
||||
@@ -242,17 +242,6 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
|
||||
--repo-id=your_dataset
|
||||
```
|
||||
|
||||
Recording, resuming, and merging aggregate quantiles from per-episode summaries, so `meta/stats.json` ends up holding a conservative envelope (`min` for `q <= 50`, `max` for `q > 50`) rather than whole-dataset quantiles. To estimate the latter, scan every episode with a running histogram:
|
||||
|
||||
```bash
|
||||
python src/lerobot/scripts/augment_dataset_quantile_stats.py \
|
||||
--repo-id=your_dataset \
|
||||
--overwrite \
|
||||
--skip-images
|
||||
```
|
||||
|
||||
`--skip-images` keeps the existing image statistics and avoids video decoding when only `STATE`/`ACTION` need recomputing, and `--root` reads a local dataset instead of the Hub. These values are histogram estimates, subject to discretization and rebinning error, so they can differ from the conservative ones — which changes MolmoAct2's normalized targets and therefore its loss scale. Statistics already saved inside an existing checkpoint are not affected.
|
||||
|
||||
Alternatively, train MolmoAct2 with mean/std normalization:
|
||||
|
||||
```bash
|
||||
|
||||
@@ -127,17 +127,6 @@ lerobot-edit-dataset \
|
||||
|
||||
Or keep the dataset as-is and pass `--policy.normalization_mapping='{"ACTION": "MEAN_STD", "STATE": "MEAN_STD", "VISUAL": "IDENTITY"}'`.
|
||||
|
||||
Recording, resuming, and merging aggregate quantiles from per-episode summaries, so `meta/stats.json` ends up holding a conservative envelope (`min` for `q <= 50`, `max` for `q > 50`) rather than whole-dataset quantiles. To estimate the latter, scan every episode with a running histogram:
|
||||
|
||||
```bash
|
||||
python src/lerobot/scripts/augment_dataset_quantile_stats.py \
|
||||
--repo-id=your_dataset \
|
||||
--overwrite \
|
||||
--skip-images
|
||||
```
|
||||
|
||||
`--skip-images` keeps the existing image statistics and avoids video decoding when only `STATE`/`ACTION` need recomputing, and `--root` reads a local dataset instead of the Hub. These values are histogram estimates, subject to discretization and rebinning error, so they can differ from the conservative ones — which changes π₀.₅'s normalized targets and therefore its loss scale. Statistics already saved inside an existing checkpoint are not affected.
|
||||
|
||||
### Training Command Example
|
||||
|
||||
The same finetune with the VLM frozen: less memory, at some cost in success rate. Swap `--dataset.repo_id` for your own dataset.
|
||||
|
||||
@@ -2,25 +2,6 @@
|
||||
|
||||
https://diffusion-policy.cs.columbia.edu
|
||||
|
||||
## Training
|
||||
|
||||
The reference implementation maintains an exponential moving average (EMA) of the policy weights during training and evaluates the EMA weights. To reproduce this behavior, enable the trainer's EMA shadow:
|
||||
|
||||
```bash
|
||||
lerobot-train \
|
||||
--policy.type=diffusion \
|
||||
--ema.enable=true \
|
||||
...
|
||||
```
|
||||
|
||||
Checkpoints then contain a directly loadable copy of the EMA weights next to the live ones, e.g. for evaluation:
|
||||
|
||||
```bash
|
||||
lerobot-eval --policy.path=outputs/train/.../checkpoints/last/pretrained_model_ema ...
|
||||
```
|
||||
|
||||
The EMA decay schedule (`--ema.inv_gamma`, `--ema.power`, ...) defaults to the reference implementation's values. For a constant decay instead of the warmup schedule (e.g. to match openpi's pi0/pi05 training), set `--ema.decay=0.99`.
|
||||
|
||||
## Citation
|
||||
|
||||
```bibtex
|
||||
|
||||
@@ -59,22 +59,6 @@ When `use_relative_actions=true`, the training script automatically:
|
||||
|
||||
---
|
||||
|
||||
## EMA of the policy weights
|
||||
|
||||
OpenPI maintains an exponential moving average of the weights during training (`ema_decay=0.99` by default) and keeps the EMA copy for inference. To reproduce this with the LeRobot trainer, enable the EMA shadow with a constant decay:
|
||||
|
||||
```bash
|
||||
python -m lerobot.scripts.lerobot_train \
|
||||
--policy.type=pi05 \
|
||||
--dataset.repo_id=your_org/your_dataset \
|
||||
--ema.enable=true \
|
||||
--ema.decay=0.99
|
||||
```
|
||||
|
||||
Checkpoints then contain a directly loadable copy of the EMA weights in `pretrained_model_ema/` next to the live ones. Note that the shadow is a full extra copy of the parameters on the GPU. Like OpenPI (which disables EMA in its LoRA configs), EMA is not supported together with PEFT adapters.
|
||||
|
||||
---
|
||||
|
||||
## Citation
|
||||
|
||||
If you use this work, please cite both **OpenPI** and the π₀.₅ paper:
|
||||
|
||||
+6
-6
@@ -408,6 +408,11 @@ ignore = [
|
||||
"T201", # Print statement found
|
||||
"T203", # Pprint statement found
|
||||
"B008", # Perform function call in argument defaults
|
||||
# D100/D104: module- and package-level docstrings. The API reference is generated from class and
|
||||
# function docstrings; a banner at the top of every file and every __init__.py would not appear on any
|
||||
# rendered page. Coverage of the things that do get rendered is enforced by interrogate instead.
|
||||
"D100",
|
||||
"D104",
|
||||
]
|
||||
|
||||
[tool.ruff.lint.per-file-ignores]
|
||||
@@ -446,8 +451,6 @@ ignore = [
|
||||
"src/lerobot/policies/**" = ["D"]
|
||||
"src/lerobot/processor/**" = ["D"]
|
||||
"src/lerobot/rewards/**" = ["D"]
|
||||
"src/lerobot/rl/**" = ["D"]
|
||||
"src/lerobot/robots/**" = ["D"]
|
||||
"src/lerobot/rollout/**" = ["D"]
|
||||
"src/lerobot/scripts/**" = ["D"]
|
||||
"src/lerobot/teleoperators/**" = ["D"]
|
||||
@@ -455,9 +458,6 @@ ignore = [
|
||||
"src/lerobot/transport/**" = ["D"]
|
||||
"src/lerobot/utils/**" = ["D"]
|
||||
"src/lerobot/lerobot_types.py" = ["D"]
|
||||
# Package root: two one-line docstring fixes land with the docstring PR.
|
||||
"src/lerobot/__init__.py" = ["D"]
|
||||
"src/lerobot/__version__.py" = ["D"]
|
||||
[tool.ruff.lint.isort]
|
||||
combine-as-imports = true
|
||||
known-first-party = ["lerobot"]
|
||||
@@ -514,7 +514,7 @@ ignore-private = false
|
||||
ignore-property-decorators = false
|
||||
ignore-module = false
|
||||
ignore-setters = false
|
||||
fail-under = 52
|
||||
fail-under = 55.5
|
||||
output-format = "term-missing"
|
||||
color = true
|
||||
paths = ["src/lerobot"]
|
||||
|
||||
@@ -14,8 +14,7 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
LeRobot -- PyTorch library for real-world robotics.
|
||||
"""LeRobot -- PyTorch library for real-world robotics.
|
||||
|
||||
Provides datasets, pretrained policies, and tools for training, evaluation,
|
||||
data collection, and robot control. Integrates with Hugging Face Hub for
|
||||
|
||||
@@ -13,7 +13,7 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""To enable `lerobot.__version__`"""
|
||||
"""To enable `lerobot.__version__`."""
|
||||
|
||||
from importlib.metadata import PackageNotFoundError, version
|
||||
|
||||
|
||||
@@ -33,10 +33,10 @@ class Camera(abc.ABC):
|
||||
- Connection/disconnection
|
||||
- Frame capture (sync/async/latest)
|
||||
|
||||
Attributes:
|
||||
fps (int | None): Configured frames per second
|
||||
width (int | None): Frame width in pixels
|
||||
height (int | None): Frame height in pixels
|
||||
**Attributes**:
|
||||
- **fps** (`int | None`) -- Configured frames per second.
|
||||
- **width** (`int | None`) -- Frame width in pixels.
|
||||
- **height** (`int | None`) -- Frame height in pixels.
|
||||
"""
|
||||
|
||||
def __init__(self, config: CameraConfig):
|
||||
|
||||
@@ -40,17 +40,20 @@ class OpenCVCameraConfig(CameraConfig):
|
||||
OpenCVCameraConfig(0, 30, 1280, 720, fourcc="YUYV") # With YUYV format
|
||||
```
|
||||
|
||||
Attributes:
|
||||
index_or_path: Either an integer representing the camera device index,
|
||||
or a Path object pointing to a video file.
|
||||
fps: Requested frames per second for the color stream.
|
||||
width: Requested frame width in pixels for the color stream.
|
||||
height: Requested frame height in pixels for the color stream.
|
||||
color_mode: Color mode for image output (RGB or BGR). Defaults to RGB.
|
||||
rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
|
||||
warmup_s: Time reading frames before returning from connect (in seconds)
|
||||
fourcc: FOURCC code for video format (e.g., "MJPG", "YUYV", "I420"). Defaults to None (auto-detect).
|
||||
backend: OpenCV backend identifier (https://docs.opencv.org/3.4/d4/d15/group__videoio__flags__base.html). Defaults to ANY.
|
||||
**Attributes**:
|
||||
- **index_or_path** (`int | Path`) -- Either an integer representing the camera device index, or a
|
||||
Path object pointing to a video file.
|
||||
- **fps** -- Requested frames per second for the color stream.
|
||||
- **width** -- Requested frame width in pixels for the color stream.
|
||||
- **height** -- Requested frame height in pixels for the color stream.
|
||||
- **color_mode** (`ColorMode`) -- Color mode for image output (RGB or BGR). Defaults to RGB.
|
||||
- **rotation** (`Cv2Rotation`) -- Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no
|
||||
rotation.
|
||||
- **warmup_s** (`int`) -- Time reading frames before returning from connect (in seconds)
|
||||
- **fourcc** (`str | None`) -- FOURCC code for video format (e.g., "MJPG", "YUYV", "I420"). Defaults
|
||||
to None (auto-detect).
|
||||
- **backend** (`Cv2Backends`) -- OpenCV backend identifier
|
||||
(https://docs.opencv.org/3.4/d4/d15/group__videoio__flags__base.html). Defaults to ANY.
|
||||
|
||||
Note:
|
||||
- Only 3-channel color output (RGB/BGR) is currently supported.
|
||||
|
||||
@@ -43,16 +43,16 @@ class Reachy2CameraConfig(CameraConfig):
|
||||
) # Left teleop camera, 640x480 @ 30FPS
|
||||
```
|
||||
|
||||
Attributes:
|
||||
name: Name of the camera device. Can be "teleop" or "depth".
|
||||
image_type: Type of image stream. For "teleop" camera, can be "left" or "right".
|
||||
For "depth" camera, can be "rgb" or "depth". (depth is not supported yet)
|
||||
fps: Requested frames per second for the color stream. Not configurable for Reachy 2 cameras.
|
||||
width: Requested frame width in pixels for the color stream.
|
||||
height: Requested frame height in pixels for the color stream.
|
||||
color_mode: Color mode for image output (RGB or BGR). Defaults to RGB.
|
||||
ip_address: IP address of the robot. Defaults to "localhost".
|
||||
port: Port number for the camera server. Defaults to 50065.
|
||||
**Attributes**:
|
||||
- **name** (`str`) -- Name of the camera device. Can be "teleop" or "depth".
|
||||
- **image_type** (`str`) -- Type of image stream. For "teleop" camera, can be "left" or "right". For
|
||||
"depth" camera, can be "rgb" or "depth". (depth is not supported yet)
|
||||
- **fps** -- Requested frames per second for the color stream. Not configurable for Reachy 2 cameras.
|
||||
- **width** -- Requested frame width in pixels for the color stream.
|
||||
- **height** -- Requested frame height in pixels for the color stream.
|
||||
- **color_mode** (`ColorMode`) -- Color mode for image output (RGB or BGR). Defaults to RGB.
|
||||
- **ip_address** (`str | None`) -- IP address of the robot. Defaults to "localhost".
|
||||
- **port** (`int`) -- Port number for the camera server. Defaults to 50065.
|
||||
|
||||
Note:
|
||||
- Only 3-channel color output (RGB/BGR) is currently supported.
|
||||
|
||||
@@ -36,27 +36,28 @@ class RealSenseCameraConfig(CameraConfig):
|
||||
RealSenseCameraConfig("0123456789", 30, 640, 480, rotation=Cv2Rotation.ROTATE_90) # With 90° rotation
|
||||
```
|
||||
|
||||
Attributes:
|
||||
fps: Requested frames per second for the color stream.
|
||||
width: Requested frame width in pixels for the color stream.
|
||||
height: Requested frame height in pixels for the color stream.
|
||||
serial_number_or_name: Unique serial number or human-readable name to identify the camera.
|
||||
color_mode: Color mode for image output (RGB or BGR). Defaults to RGB.
|
||||
use_rgb: Whether to enable the color stream. Defaults to True.
|
||||
use_depth: Whether to enable depth stream. Defaults to False.
|
||||
rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
|
||||
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).
|
||||
**Attributes**:
|
||||
- **fps** -- Requested frames per second for the color stream.
|
||||
- **width** -- Requested frame width in pixels for the color stream.
|
||||
- **height** -- Requested frame height in pixels for the color stream.
|
||||
- **serial_number_or_name** (`str`) -- Unique serial number or human-readable name to identify the
|
||||
camera.
|
||||
- **color_mode** (`ColorMode`) -- Color mode for image output (RGB or BGR). Defaults to RGB.
|
||||
- **use_rgb** (`bool`) -- Whether to enable the color stream. Defaults to True.
|
||||
- **use_depth** (`bool`) -- Whether to enable depth stream. Defaults to False.
|
||||
- **rotation** (`Cv2Rotation`) -- Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no
|
||||
rotation.
|
||||
- **warmup_s** (`int`) -- Time reading frames before returning from connect (in seconds)
|
||||
- **exposure** (`int | None`) -- 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** (`int | None`) -- 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** (`int | None`) -- 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:
|
||||
- Either name or serial_number must be specified.
|
||||
|
||||
@@ -22,7 +22,7 @@ Import them directly: ``from lerobot.configs.train import TrainPipelineConfig``
|
||||
"""
|
||||
|
||||
from .dataset import DatasetRecordConfig
|
||||
from .default import DatasetConfig, EMAConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
|
||||
from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
|
||||
from .policies import PreTrainedConfig
|
||||
from .recipe import MessageTurn, TrainingRecipe, load_recipe
|
||||
from .types import (
|
||||
@@ -57,7 +57,6 @@ __all__ = [
|
||||
# Config classes
|
||||
"DatasetRecordConfig",
|
||||
"DatasetConfig",
|
||||
"EMAConfig",
|
||||
"EvalConfig",
|
||||
"JobConfig",
|
||||
"MessageTurn",
|
||||
|
||||
@@ -139,59 +139,6 @@ class EvalConfig:
|
||||
return min(by_cpu, self.n_episodes, 64)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EMAConfig:
|
||||
"""Exponential moving average (EMA) of the policy weights.
|
||||
|
||||
Standard practice for diffusion-style policies (Chi et al. 2023, "Diffusion Policy", section V.D):
|
||||
the reference implementation enables it in every config and evaluates the EMA weights. Off by
|
||||
default here because it keeps a second full copy of the parameters in memory.
|
||||
|
||||
The decay follows the warmup schedule from diffusers' `EMAModel`:
|
||||
`decay_t = 1 - (1 + t / inv_gamma) ** -power`, clamped to `[min_decay, max_decay]`.
|
||||
The defaults mirror the reference implementation. Alternatively, set `decay` for a constant
|
||||
decay at every step, as used by openpi for pi0/pi05 (`ema_decay=0.99`).
|
||||
"""
|
||||
|
||||
enable: bool = False
|
||||
# Constant decay coefficient (openpi-style, e.g. 0.99 for pi0/pi05). When set, the warmup
|
||||
# schedule below is bypassed and the shadow uses this decay at every step.
|
||||
decay: float | None = None
|
||||
# Number of optimizer steps during which the shadow stays a hard copy of the live weights.
|
||||
update_after_step: int = 0
|
||||
# Warmup schedule parameters (see class docstring).
|
||||
inv_gamma: float = 1.0
|
||||
power: float = 0.75
|
||||
min_decay: float = 0.0
|
||||
max_decay: float = 0.9999
|
||||
# Evaluate the EMA weights (instead of the live ones) during periodic env eval.
|
||||
# Offline eval-loss (--eval_steps) always uses the live weights: it runs on every rank
|
||||
# while the EMA shadow only lives on the main process.
|
||||
use_for_eval: bool = True
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not (0.0 <= self.min_decay <= self.max_decay <= 1.0):
|
||||
raise ValueError(
|
||||
"Expected 0 <= ema.min_decay <= ema.max_decay <= 1, got "
|
||||
f"min_decay={self.min_decay} and max_decay={self.max_decay}."
|
||||
)
|
||||
if self.inv_gamma <= 0:
|
||||
raise ValueError(f"ema.inv_gamma must be positive, got {self.inv_gamma}.")
|
||||
if self.power <= 0:
|
||||
raise ValueError(f"ema.power must be positive, got {self.power}.")
|
||||
if self.update_after_step < 0:
|
||||
raise ValueError(f"ema.update_after_step must be >= 0, got {self.update_after_step}.")
|
||||
if self.decay is not None:
|
||||
if not 0.0 <= self.decay <= 1.0:
|
||||
raise ValueError(f"ema.decay must be in [0, 1], got {self.decay}.")
|
||||
# Keep the literals in sync with the field defaults above.
|
||||
if self.min_decay != 0.0 or self.max_decay != 0.9999:
|
||||
raise ValueError(
|
||||
"ema.decay (constant decay) and ema.min_decay/ema.max_decay (schedule clamp) are "
|
||||
"mutually exclusive: set one or the other."
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PeftConfig:
|
||||
# PEFT offers many fine-tuning methods, layer adapters being the most common and currently also the most
|
||||
|
||||
@@ -35,7 +35,7 @@ from lerobot.utils.hub import HubMixin, find_latest_hub_checkpoint
|
||||
from lerobot.utils.sample_weighting import SampleWeightingConfig
|
||||
|
||||
from . import parser
|
||||
from .default import DatasetConfig, EMAConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
|
||||
from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
|
||||
from .policies import PreTrainedConfig
|
||||
from .rewards import RewardModelConfig
|
||||
|
||||
@@ -163,8 +163,6 @@ class TrainPipelineConfig(HubMixin):
|
||||
# FSDP/DDP tuning knobs, compile & activation-checkpointing placeholders.
|
||||
accelerator: AcceleratorConfig = field(default_factory=AcceleratorConfig)
|
||||
eval: EvalConfig = field(default_factory=EvalConfig)
|
||||
# Maintain an EMA shadow of the policy weights during training (see EMAConfig).
|
||||
ema: EMAConfig = field(default_factory=EMAConfig)
|
||||
wandb: WandBConfig = field(default_factory=WandBConfig)
|
||||
peft: PeftConfig | None = None
|
||||
|
||||
|
||||
@@ -613,15 +613,8 @@ def aggregate_feature_stats(stats_ft_list: list[dict[str, dict]]) -> dict[str, d
|
||||
for q_key in quantile_keys:
|
||||
if all(q_key in s for s in stats_ft_list):
|
||||
quantile_values = np.stack([s[q_key] for s in stats_ft_list])
|
||||
# Exact global quantiles cannot be recovered from quantile summaries.
|
||||
# Keep a conservative envelope of the available estimates: min
|
||||
# for lower quantiles and max for upper quantiles. The resulting
|
||||
# values are bounds across the inputs, not global quantile estimates.
|
||||
q_percent = int(q_key[1:])
|
||||
if q_percent <= 50:
|
||||
aggregated[q_key] = np.min(quantile_values, axis=0)
|
||||
else:
|
||||
aggregated[q_key] = np.max(quantile_values, axis=0)
|
||||
weighted_quantiles = quantile_values * counts
|
||||
aggregated[q_key] = weighted_quantiles.sum(axis=0) / total_count
|
||||
|
||||
return aggregated
|
||||
|
||||
|
||||
@@ -33,7 +33,11 @@ if TYPE_CHECKING:
|
||||
from lerobot.policies.pretrained import PreTrainedPolicy
|
||||
|
||||
# Env vars through which `accelerate launch --config_file` (or a stray shell) would configure
|
||||
# accelerate behind the config system's back, making train_config.json lie about what ran.
|
||||
# accelerate behind the config system's back. Plugin `__post_init__`s read these silently as
|
||||
# field fallbacks (ACCELERATE_DYNAMO_* enables torch.compile through the default
|
||||
# TorchDynamoPlugin; ACCELERATE_GRADIENT_ACCUMULATION_STEPS overrides the explicitly passed
|
||||
# value inside Accelerator.__init__), which would make train_config.json lie about what ran.
|
||||
_ACCELERATE_ENV_PREFIXES = ("FSDP_", "PARALLELISM_CONFIG_", "ACCELERATE_DYNAMO_")
|
||||
_ACCELERATE_ENV_VARS = (
|
||||
"ACCELERATE_USE_FSDP",
|
||||
"ACCELERATE_USE_PARALLELISM_CONFIG",
|
||||
@@ -55,7 +59,11 @@ def guard_against_env_interference() -> None:
|
||||
"""
|
||||
if os.environ.get(_ENV_OVERRIDE):
|
||||
return
|
||||
offending = sorted(name for name in _ACCELERATE_ENV_VARS if name in os.environ)
|
||||
offending = sorted(
|
||||
name
|
||||
for name in os.environ
|
||||
if name in _ACCELERATE_ENV_VARS or name.startswith(_ACCELERATE_ENV_PREFIXES)
|
||||
)
|
||||
if offending:
|
||||
raise RuntimeError(
|
||||
f"Accelerate-configuring environment variables are set: {', '.join(offending)}. "
|
||||
|
||||
@@ -314,11 +314,16 @@ class SerialMotorsBus(MotorsBusBase):
|
||||
To find the port, you can run our utility script:
|
||||
```bash
|
||||
lerobot-find-port.py
|
||||
>>> Finding all available ports for the MotorsBus.
|
||||
>>> ["/dev/tty.usbmodem575E0032081", "/dev/tty.usbmodem575E0031751"]
|
||||
>>> Remove the usb cable from your MotorsBus and press Enter when done.
|
||||
>>> The port of this MotorsBus is /dev/tty.usbmodem575E0031751.
|
||||
>>> Reconnect the usb cable.
|
||||
```
|
||||
|
||||
which prints:
|
||||
|
||||
```
|
||||
Finding all available ports for the MotorsBus.
|
||||
["/dev/tty.usbmodem575E0032081", "/dev/tty.usbmodem575E0031751"]
|
||||
Remove the usb cable from your MotorsBus and press Enter when done.
|
||||
The port of this MotorsBus is /dev/tty.usbmodem575E0031751.
|
||||
Reconnect the usb cable.
|
||||
```
|
||||
|
||||
Example of usage for 1 Feetech sts3215 motor connected to the bus:
|
||||
@@ -595,7 +600,7 @@ class SerialMotorsBus(MotorsBusBase):
|
||||
ID, and finally programs the bus' default baud-rate.
|
||||
|
||||
Args:
|
||||
motor (str): Key of the motor in :pyattr:`motors`.
|
||||
motor (str): Key of the motor in `motors`.
|
||||
initial_baudrate (int | None, optional): Current baud-rate (skips scanning when provided).
|
||||
Defaults to None.
|
||||
initial_id (int | None, optional): Current ID (skips scanning when provided). Defaults to None.
|
||||
@@ -666,7 +671,7 @@ class SerialMotorsBus(MotorsBusBase):
|
||||
"""Enable torque on selected motors.
|
||||
|
||||
Args:
|
||||
motors (int | str | list[str] | None, optional): Same semantics as :pymeth:`disable_torque`.
|
||||
motors (int | str | list[str] | None, optional): Same semantics as [`~motors.motors_bus.MotorsBus.disable_torque`].
|
||||
Defaults to `None`.
|
||||
num_retry (int, optional): Number of additional retry attempts on communication failure.
|
||||
Defaults to 0.
|
||||
@@ -679,10 +684,12 @@ class SerialMotorsBus(MotorsBusBase):
|
||||
|
||||
This helper is useful to temporarily disable torque when configuring motors.
|
||||
|
||||
Examples:
|
||||
>>> with bus.torque_disabled():
|
||||
Example:
|
||||
```python
|
||||
>>> with bus.torque_disabled(): # doctest: +SKIP
|
||||
... # Safe operations here
|
||||
... pass
|
||||
```
|
||||
"""
|
||||
self.disable_torque(motors)
|
||||
try:
|
||||
@@ -695,7 +702,7 @@ class SerialMotorsBus(MotorsBusBase):
|
||||
|
||||
Args:
|
||||
timeout_ms (int | None, optional): Timeout in *milliseconds*. If `None` (default) the method falls
|
||||
back to :pyattr:`default_timeout`.
|
||||
back to `default_timeout`.
|
||||
"""
|
||||
timeout_ms = timeout_ms if timeout_ms is not None else self.default_timeout
|
||||
self.port_handler.setPacketTimeoutMillis(timeout_ms)
|
||||
@@ -746,8 +753,8 @@ class SerialMotorsBus(MotorsBusBase):
|
||||
|
||||
Args:
|
||||
calibration_dict (dict[str, MotorCalibration]): Calibration obtained from
|
||||
:pymeth:`read_calibration` or crafted by the user.
|
||||
cache (bool, optional): Save the calibration to :pyattr:`calibration`. Defaults to True.
|
||||
[`~motors.motors_bus.MotorsBus.read_calibration`] or crafted by the user.
|
||||
cache (bool, optional): Save the calibration to `calibration`. Defaults to True.
|
||||
"""
|
||||
pass
|
||||
|
||||
@@ -755,7 +762,7 @@ class SerialMotorsBus(MotorsBusBase):
|
||||
"""Restore factory calibration for the selected motors.
|
||||
|
||||
Homing offset is set to ``0`` and min/max position limits are set to the full usable range.
|
||||
The in-memory :pyattr:`calibration` is cleared.
|
||||
The in-memory `calibration` is cleared.
|
||||
|
||||
Args:
|
||||
motors (NameOrID | Sequence[NameOrID] | None, optional): Selection of motors. `None` (default)
|
||||
@@ -1069,9 +1076,9 @@ class SerialMotorsBus(MotorsBusBase):
|
||||
) -> None:
|
||||
"""Write a value to a single motor's register.
|
||||
|
||||
Contrary to :pymeth:`sync_write`, this expects a response status packet emitted by the motor, which
|
||||
Contrary to [`~motors.motors_bus.MotorsBus.sync_write`], this expects a response status packet emitted by the motor, which
|
||||
provides a guarantee that the value was written to the register successfully. In consequence, it is
|
||||
slower than :pymeth:`sync_write` but it is more reliable. It should typically be used when configuring
|
||||
slower than [`~motors.motors_bus.MotorsBus.sync_write`] but it is more reliable. It should typically be used when configuring
|
||||
motors.
|
||||
|
||||
Args:
|
||||
@@ -1228,8 +1235,8 @@ class SerialMotorsBus(MotorsBusBase):
|
||||
) -> None:
|
||||
"""Write the same register on multiple motors.
|
||||
|
||||
Contrary to :pymeth:`write`, this *does not* expects a response status packet emitted by the motor, which
|
||||
can allow for lost packets. It is faster than :pymeth:`write` and should typically be used when
|
||||
Contrary to [`~motors.motors_bus.MotorsBus.write`], this *does not* expects a response status packet emitted by the motor, which
|
||||
can allow for lost packets. It is faster than [`~motors.motors_bus.MotorsBus.write`] and should typically be used when
|
||||
frequency matters and losing some packets is acceptable (e.g. teleoperation loops).
|
||||
|
||||
Args:
|
||||
|
||||
@@ -131,12 +131,16 @@ class ProcessorConfigKwargs(TypedDict, total=False):
|
||||
This provides type hints for the optional arguments passed to `make_pre_post_processors`,
|
||||
improving code clarity and enabling static analysis.
|
||||
|
||||
Attributes:
|
||||
preprocessor_config_filename: The filename for the preprocessor configuration.
|
||||
postprocessor_config_filename: The filename for the postprocessor configuration.
|
||||
preprocessor_overrides: A dictionary of overrides for the preprocessor configuration.
|
||||
postprocessor_overrides: A dictionary of overrides for the postprocessor configuration.
|
||||
dataset_stats: Dataset statistics for normalization.
|
||||
**Attributes**:
|
||||
- **preprocessor_config_filename** (`str | None`) -- The filename for the preprocessor configuration.
|
||||
- **postprocessor_config_filename** (`str | None`) -- The filename for the postprocessor
|
||||
configuration.
|
||||
- **preprocessor_overrides** (`dict[str, Any] | None`) -- A dictionary of overrides for the
|
||||
preprocessor configuration.
|
||||
- **postprocessor_overrides** (`dict[str, Any] | None`) -- A dictionary of overrides for the
|
||||
postprocessor configuration.
|
||||
- **dataset_stats** (`dict[str, dict[str, torch.Tensor]] | None`) -- Dataset statistics for
|
||||
normalization.
|
||||
"""
|
||||
|
||||
preprocessor_config_filename: str | None
|
||||
|
||||
@@ -46,10 +46,11 @@ class ActionQueue:
|
||||
Args:
|
||||
cfg (RTCConfig): Configuration for Real-Time Chunking behavior.
|
||||
|
||||
Attributes:
|
||||
queue (Tensor | None): Processed actions for robot rollout (time_steps, action_dim).
|
||||
original_queue (Tensor | None): Original actions for RTC computation (time_steps, action_dim).
|
||||
last_index (int): Current consumption index in the queue.
|
||||
**Attributes**:
|
||||
- **queue** (`Tensor | None`) -- Processed actions for robot rollout (time_steps, action_dim).
|
||||
- **original_queue** (`Tensor | None`) -- Original actions for RTC computation (time_steps,
|
||||
action_dim).
|
||||
- **last_index** (`int`) -- Current consumption index in the queue.
|
||||
"""
|
||||
|
||||
def __init__(self, cfg: RTCConfig):
|
||||
|
||||
@@ -27,19 +27,19 @@ from torch import Tensor
|
||||
class DebugStep:
|
||||
"""Container for debug information from a single denoising step.
|
||||
|
||||
Attributes:
|
||||
step_idx (int): Step index/counter.
|
||||
x_t (Tensor | None): Current latent/state tensor.
|
||||
v_t (Tensor | None): Velocity from denoiser.
|
||||
x1_t (Tensor | None): Denoised prediction (x_t - time * v_t).
|
||||
correction (Tensor | None): Correction gradient tensor.
|
||||
err (Tensor | None): Weighted error term.
|
||||
weights (Tensor | None): Prefix attention weights.
|
||||
guidance_weight (float | Tensor | None): Applied guidance weight.
|
||||
time (float | Tensor | None): Time parameter.
|
||||
inference_delay (int | None): Inference delay parameter.
|
||||
execution_horizon (int | None): Execution horizon parameter.
|
||||
metadata (dict[str, Any]): Additional metadata.
|
||||
**Attributes**:
|
||||
- **step_idx** (`int`) -- Step index/counter.
|
||||
- **x_t** (`Tensor | None`) -- Current latent/state tensor.
|
||||
- **v_t** (`Tensor | None`) -- Velocity from denoiser.
|
||||
- **x1_t** (`Tensor | None`) -- Denoised prediction (x_t - time * v_t).
|
||||
- **correction** (`Tensor | None`) -- Correction gradient tensor.
|
||||
- **err** (`Tensor | None`) -- Weighted error term.
|
||||
- **weights** (`Tensor | None`) -- Prefix attention weights.
|
||||
- **guidance_weight** (`float | Tensor | None`) -- Applied guidance weight.
|
||||
- **time** (`float | Tensor | None`) -- Time parameter.
|
||||
- **inference_delay** (`int | None`) -- Inference delay parameter.
|
||||
- **execution_horizon** (`int | None`) -- Execution horizon parameter.
|
||||
- **metadata** (`dict[str, Any]`) -- Additional metadata.
|
||||
"""
|
||||
|
||||
step_idx: int = 0
|
||||
|
||||
@@ -217,10 +217,12 @@ class AddBatchDimensionProcessorStep(ProcessorStep):
|
||||
This step combines individual processors for actions, observations, and complementary data
|
||||
to create a batched transition (batch size 1) from a single-instance transition.
|
||||
|
||||
Attributes:
|
||||
to_batch_action_processor: Processor for the action component.
|
||||
to_batch_observation_processor: Processor for the observation component.
|
||||
to_batch_complementary_data_processor: Processor for the complementary data component.
|
||||
**Attributes**:
|
||||
- **to_batch_action_processor** (`AddBatchDimensionActionStep`) -- Processor for the action component.
|
||||
- **to_batch_observation_processor** (`AddBatchDimensionObservationStep`) -- Processor for the
|
||||
observation component.
|
||||
- **to_batch_complementary_data_processor** (`AddBatchDimensionComplementaryDataStep`) -- Processor
|
||||
for the complementary data component.
|
||||
"""
|
||||
|
||||
to_batch_action_processor: AddBatchDimensionActionStep = field(
|
||||
|
||||
@@ -32,9 +32,8 @@ class MapTensorToDeltaActionDictStep(ActionProcessorStep):
|
||||
It decomposes the vector into named components for delta movements of the
|
||||
end-effector (x, y, z) and optionally the gripper.
|
||||
|
||||
Attributes:
|
||||
use_gripper: If True, assumes the 4th element of the tensor is the
|
||||
gripper action.
|
||||
**Attributes**:
|
||||
- **use_gripper** (`bool`) -- If True, assumes the 4th element of the tensor is the gripper action.
|
||||
"""
|
||||
|
||||
use_gripper: bool = True
|
||||
@@ -81,10 +80,10 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
|
||||
into a target action format that includes an "enabled" flag and target
|
||||
end-effector positions. It also handles scaling and noise filtering.
|
||||
|
||||
Attributes:
|
||||
position_scale: A factor to scale the delta position inputs.
|
||||
noise_threshold: The magnitude below which delta inputs are considered noise
|
||||
and do not trigger an "enabled" state.
|
||||
**Attributes**:
|
||||
- **position_scale** (`float`) -- A factor to scale the delta position inputs.
|
||||
- **noise_threshold** (`float`) -- The magnitude below which delta inputs are considered noise and do
|
||||
not trigger an "enabled" state.
|
||||
"""
|
||||
|
||||
# Scale factors for delta movements
|
||||
|
||||
@@ -40,10 +40,10 @@ class DeviceProcessorStep(ProcessorStep):
|
||||
|
||||
This is crucial for preparing data for model training or inference on hardware like GPUs.
|
||||
|
||||
Attributes:
|
||||
device: The target device for tensors (e.g., "cpu", "cuda", "cuda:0").
|
||||
float_dtype: The target floating-point dtype as a string (e.g., "float32", "float16", "bfloat16").
|
||||
If None, the dtype is not changed.
|
||||
**Attributes**:
|
||||
- **device** (`str`) -- The target device for tensors (e.g., "cpu", "cuda", "cuda:0").
|
||||
- **float_dtype** (`str | None`) -- The target floating-point dtype as a string (e.g., "float32",
|
||||
"float16", "bfloat16"). If None, the dtype is not changed.
|
||||
"""
|
||||
|
||||
device: str = "cpu"
|
||||
|
||||
@@ -33,10 +33,9 @@ class Torch2NumpyActionProcessorStep(ActionProcessorStep):
|
||||
This step is useful when the output of a policy (typically a torch.Tensor)
|
||||
needs to be passed to an environment or component that expects a NumPy array.
|
||||
|
||||
Attributes:
|
||||
squeeze_batch_dim: If True, removes the first dimension of the array
|
||||
if it is of size 1. This is useful for converting a
|
||||
batched action of size (1, D) to a single action of size (D,).
|
||||
**Attributes**:
|
||||
- **squeeze_batch_dim** (`bool`) -- If True, removes the first dimension of the array if it is of size
|
||||
1. This is useful for converting a batched action of size (1, D) to a single action of size (D,).
|
||||
"""
|
||||
|
||||
squeeze_batch_dim: bool = True
|
||||
|
||||
@@ -101,8 +101,8 @@ class AddTeleopActionAsComplimentaryDataStep(ComplementaryDataProcessorStep):
|
||||
be available to downstream processors, for example, to override a policy's action
|
||||
during an intervention.
|
||||
|
||||
Attributes:
|
||||
teleop_device: The teleoperator instance to get the action from.
|
||||
**Attributes**:
|
||||
- **teleop_device** (`Teleoperator`) -- The teleoperator instance to get the action from.
|
||||
"""
|
||||
|
||||
teleop_device: "Teleoperator"
|
||||
@@ -137,9 +137,9 @@ class AddTeleopEventsAsInfoStep(InfoProcessorStep):
|
||||
This step extracts control events from teleoperators that support event-based
|
||||
interaction, making these signals available to other parts of the system.
|
||||
|
||||
Attributes:
|
||||
teleop_device: An instance of a teleoperator that implements the
|
||||
`HasTeleopEvents` protocol.
|
||||
**Attributes**:
|
||||
- **teleop_device** (`TeleopWithEvents`) -- An instance of a teleoperator that implements the
|
||||
`HasTeleopEvents` protocol.
|
||||
"""
|
||||
|
||||
teleop_device: TeleopWithEvents
|
||||
@@ -180,10 +180,10 @@ class ImageCropResizeProcessorStep(ObservationProcessorStep):
|
||||
the specified transformations. It handles device placement, moving tensors to the
|
||||
CPU if necessary for operations not supported on certain accelerators like MPS.
|
||||
|
||||
Attributes:
|
||||
crop_params_dict: A dictionary mapping image keys to cropping parameters
|
||||
(top, left, height, width).
|
||||
resize_size: A tuple (height, width) to resize all images to.
|
||||
**Attributes**:
|
||||
- **crop_params_dict** (`dict[str, tuple[int, int, int, int]] | None`) -- A dictionary mapping image
|
||||
keys to cropping parameters (top, left, height, width).
|
||||
- **resize_size** (`tuple[int, int] | None`) -- A tuple (height, width) to resize all images to.
|
||||
"""
|
||||
|
||||
crop_params_dict: dict[str, tuple[int, int, int, int]] | None = None
|
||||
@@ -267,9 +267,9 @@ class TimeLimitProcessorStep(TruncatedProcessorStep):
|
||||
"""
|
||||
Tracks episode steps and enforces a time limit by truncating the episode.
|
||||
|
||||
Attributes:
|
||||
max_episode_steps: The maximum number of steps allowed per episode.
|
||||
current_step: The current step count for the active episode.
|
||||
**Attributes**:
|
||||
- **max_episode_steps** (`int`) -- The maximum number of steps allowed per episode.
|
||||
- **current_step** (`int`) -- The current step count for the active episode.
|
||||
"""
|
||||
|
||||
max_episode_steps: int
|
||||
@@ -358,11 +358,11 @@ class GripperPenaltyProcessorStep(ProcessorStep):
|
||||
This discourages gripper oscillation while leaving "stay" and saturating-further
|
||||
commands unpenalized.
|
||||
|
||||
Attributes:
|
||||
penalty: The negative reward value to apply.
|
||||
max_gripper_pos: The maximum position value for the gripper, used for normalization.
|
||||
open_threshold: Normalized state below which the gripper is considered "open".
|
||||
closed_threshold: Normalized state above which the gripper is considered "closed".
|
||||
**Attributes**:
|
||||
- **penalty** (`float`) -- The negative reward value to apply.
|
||||
- **max_gripper_pos** (`float`) -- The maximum position value for the gripper, used for normalization.
|
||||
- **open_threshold** (`float`) -- Normalized state below which the gripper is considered "open".
|
||||
- **closed_threshold** (`float`) -- Normalized state above which the gripper is considered "closed".
|
||||
"""
|
||||
|
||||
penalty: float = -0.02
|
||||
@@ -456,10 +456,10 @@ class InterventionActionProcessorStep(ProcessorStep):
|
||||
this step replaces the policy's action with the human's teleoperated action.
|
||||
It also processes signals to terminate the episode or flag success.
|
||||
|
||||
Attributes:
|
||||
use_gripper: Whether to include the gripper in the teleoperated action.
|
||||
terminate_on_success: If True, automatically sets the `done` flag when a
|
||||
`success` event is received.
|
||||
**Attributes**:
|
||||
- **use_gripper** (`bool`) -- Whether to include the gripper in the teleoperated action.
|
||||
- **terminate_on_success** (`bool`) -- If True, automatically sets the `done` flag when a `success`
|
||||
event is received.
|
||||
"""
|
||||
|
||||
use_gripper: bool = False
|
||||
@@ -557,13 +557,13 @@ class RewardClassifierProcessorStep(ProcessorStep):
|
||||
This step uses a model to determine if the current state is successful, updating
|
||||
the reward and potentially terminating the episode.
|
||||
|
||||
Attributes:
|
||||
pretrained_path: Path to the pretrained reward classifier model.
|
||||
device: The device to run the classifier on.
|
||||
success_threshold: The probability threshold to consider a prediction as successful.
|
||||
success_reward: The reward value to assign on success.
|
||||
terminate_on_success: If True, terminates the episode upon successful classification.
|
||||
reward_classifier: The loaded classifier model instance.
|
||||
**Attributes**:
|
||||
- **pretrained_path** (`str | None`) -- Path to the pretrained reward classifier model.
|
||||
- **device** (`str`) -- The device to run the classifier on.
|
||||
- **success_threshold** (`float`) -- The probability threshold to consider a prediction as successful.
|
||||
- **success_reward** (`float`) -- The reward value to assign on success.
|
||||
- **terminate_on_success** (`bool`) -- If True, terminates the episode upon successful classification.
|
||||
- **reward_classifier** (`Any`) -- The loaded classifier model instance.
|
||||
"""
|
||||
|
||||
pretrained_path: str | None = None
|
||||
|
||||
@@ -71,22 +71,23 @@ class _NormalizationMixin:
|
||||
)
|
||||
```
|
||||
|
||||
Attributes:
|
||||
features: A dictionary mapping feature names to `PolicyFeature` objects, defining
|
||||
the data structure to be processed.
|
||||
norm_map: A dictionary mapping `FeatureType` to `NormalizationMode`, specifying
|
||||
which normalization method to use for each type of feature.
|
||||
stats: A dictionary containing the normalization statistics (e.g., mean, std,
|
||||
min, max) for each feature.
|
||||
device: The PyTorch device on which to store and perform tensor operations.
|
||||
eps: A small epsilon value to prevent division by zero in normalization
|
||||
calculations.
|
||||
normalize_observation_keys: An optional set of keys to selectively apply
|
||||
normalization to specific observation features.
|
||||
_tensor_stats: An internal dictionary holding the normalization statistics as
|
||||
PyTorch tensors.
|
||||
_stats_explicitly_provided: Internal flag tracking whether stats were explicitly
|
||||
provided during construction (used for override preservation).
|
||||
**Attributes**:
|
||||
- **features** (`dict[str, PolicyFeature]`) -- A dictionary mapping feature names to `PolicyFeature`
|
||||
objects, defining the data structure to be processed.
|
||||
- **norm_map** (`dict[FeatureType, NormalizationMode]`) -- A dictionary mapping `FeatureType` to
|
||||
`NormalizationMode`, specifying which normalization method to use for each type of feature.
|
||||
- **stats** (`dict[str, dict[str, Any]] | None`) -- A dictionary containing the normalization
|
||||
statistics (e.g., mean, std, min, max) for each feature.
|
||||
- **device** (`torch.device | str | None`) -- The PyTorch device on which to store and perform tensor
|
||||
operations.
|
||||
- **eps** (`float`) -- A small epsilon value to prevent division by zero in normalization
|
||||
calculations.
|
||||
- **normalize_observation_keys** (`set[str] | None`) -- An optional set of keys to selectively apply
|
||||
normalization to specific observation features.
|
||||
- **_tensor_stats** (`dict[str, dict[str, Tensor]]`) -- An internal dictionary holding the
|
||||
normalization statistics as PyTorch tensors.
|
||||
- **_stats_explicitly_provided** (`bool`) -- Internal flag tracking whether stats were explicitly
|
||||
provided during construction (used for override preservation).
|
||||
"""
|
||||
|
||||
features: dict[str, PolicyFeature]
|
||||
|
||||
@@ -269,13 +269,18 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
||||
data processing workflow. It's generic, allowing for custom input and output types,
|
||||
which are handled by the `to_transition` and `to_output` converters.
|
||||
|
||||
Attributes:
|
||||
steps: A sequence of `ProcessorStep` objects that make up the pipeline.
|
||||
name: A descriptive name for the pipeline.
|
||||
to_transition: A function to convert raw input data into the standardized `EnvTransition` format.
|
||||
to_output: A function to convert the final `EnvTransition` into the desired output format.
|
||||
before_step_hooks: A list of functions to be called before each step is executed.
|
||||
after_step_hooks: A list of functions to be called after each step is executed.
|
||||
**Attributes**:
|
||||
- **steps** (`Sequence[ProcessorStep]`) -- A sequence of `ProcessorStep` objects that make up the
|
||||
pipeline.
|
||||
- **name** (`str`) -- A descriptive name for the pipeline.
|
||||
- **to_transition** (`Callable[[TInput], EnvTransition]`) -- A function to convert raw input data into
|
||||
the standardized `EnvTransition` format.
|
||||
- **to_output** (`Callable[[EnvTransition], TOutput]`) -- A function to convert the final
|
||||
`EnvTransition` into the desired output format.
|
||||
- **before_step_hooks** (`list[Callable[[int, EnvTransition], None]]`) -- A list of functions to be
|
||||
called before each step is executed.
|
||||
- **after_step_hooks** (`list[Callable[[int, EnvTransition], None]]`) -- A list of functions to be
|
||||
called after each step is executed.
|
||||
"""
|
||||
|
||||
steps: Sequence[ProcessorStep] = field(default_factory=list)
|
||||
|
||||
@@ -91,11 +91,11 @@ class RelativeActionsProcessorStep(ProcessorStep):
|
||||
Caches the last seen state so a paired AbsoluteActionsProcessorStep can reverse
|
||||
the conversion during postprocessing.
|
||||
|
||||
Attributes:
|
||||
enabled: Whether to apply the relative conversion.
|
||||
exclude_joints: Joint names to keep absolute (not converted to relative).
|
||||
action_names: Action dimension names from dataset metadata, used to build
|
||||
the mask from exclude_joints. If None, all dims are converted.
|
||||
**Attributes**:
|
||||
- **enabled** (`bool`) -- Whether to apply the relative conversion.
|
||||
- **exclude_joints** (`list[str]`) -- Joint names to keep absolute (not converted to relative).
|
||||
- **action_names** (`list[str] | None`) -- Action dimension names from dataset metadata, used to build
|
||||
the mask from exclude_joints. If None, all dims are converted.
|
||||
"""
|
||||
|
||||
enabled: bool = False
|
||||
@@ -168,9 +168,10 @@ class AbsoluteActionsProcessorStep(ProcessorStep):
|
||||
predicted relative offsets are converted back to absolute positions for execution.
|
||||
Reads the cached state from its paired RelativeActionsProcessorStep.
|
||||
|
||||
Attributes:
|
||||
enabled: Whether to apply the absolute conversion.
|
||||
relative_step: Reference to the paired RelativeActionsProcessorStep that caches state.
|
||||
**Attributes**:
|
||||
- **enabled** (`bool`) -- Whether to apply the absolute conversion.
|
||||
- **relative_step** (`RelativeActionsProcessorStep | None`) -- Reference to the paired
|
||||
RelativeActionsProcessorStep that caches state.
|
||||
"""
|
||||
|
||||
enabled: bool = False
|
||||
|
||||
@@ -32,10 +32,9 @@ class RenameObservationsProcessorStep(ObservationProcessorStep):
|
||||
from an environment's format to the format expected by a LeRobot policy or
|
||||
other downstream components.
|
||||
|
||||
Attributes:
|
||||
rename_map: A dictionary mapping from old key names to new key names.
|
||||
Keys present in an observation that are not in this map will
|
||||
be kept with their original names.
|
||||
**Attributes**:
|
||||
- **rename_map** (`dict[str, str]`) -- A dictionary mapping from old key names to new key names. Keys
|
||||
present in an observation that are not in this map will be kept with their original names.
|
||||
"""
|
||||
|
||||
rename_map: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
@@ -65,15 +65,17 @@ class TokenizerProcessorStep(ObservationProcessorStep):
|
||||
|
||||
Requires the `transformers` library to be installed.
|
||||
|
||||
Attributes:
|
||||
tokenizer_name: The name of a pretrained tokenizer from the Hugging Face Hub (e.g., "bert-base-uncased").
|
||||
tokenizer: A pre-initialized tokenizer object. If provided, `tokenizer_name` is ignored.
|
||||
max_length: The maximum length to pad or truncate sequences to.
|
||||
task_key: The key in `complementary_data` where the task string is stored.
|
||||
padding_side: The side to pad on ('left' or 'right').
|
||||
padding: The padding strategy ('max_length', 'longest', etc.).
|
||||
truncation: Whether to truncate sequences longer than `max_length`.
|
||||
input_tokenizer: The internal tokenizer instance, loaded during initialization.
|
||||
**Attributes**:
|
||||
- **tokenizer_name** (`str | None`) -- The name of a pretrained tokenizer from the Hugging Face Hub
|
||||
(e.g., "bert-base-uncased").
|
||||
- **tokenizer** (`Any | None`) -- A pre-initialized tokenizer object. If provided, `tokenizer_name` is
|
||||
ignored.
|
||||
- **max_length** (`int`) -- The maximum length to pad or truncate sequences to.
|
||||
- **task_key** (`str`) -- The key in `complementary_data` where the task string is stored.
|
||||
- **padding_side** (`str`) -- The side to pad on ('left' or 'right').
|
||||
- **padding** (`str`) -- The padding strategy ('max_length', 'longest', etc.).
|
||||
- **truncation** (`bool`) -- Whether to truncate sequences longer than `max_length`.
|
||||
- **input_tokenizer** (`Any`) -- The internal tokenizer instance, loaded during initialization.
|
||||
"""
|
||||
|
||||
tokenizer_name: str | None = None
|
||||
@@ -346,12 +348,17 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
|
||||
|
||||
Requires the `transformers` library to be installed.
|
||||
|
||||
Attributes:
|
||||
tokenizer_name: The name of a pretrained processor from the Hugging Face Hub (e.g., "lerobot/fast-action-tokenizer").
|
||||
tokenizer: A pre-initialized processor/tokenizer object. If provided, `tokenizer_name` is ignored.
|
||||
trust_remote_code: Whether to trust remote code when loading the tokenizer (required for some tokenizers).
|
||||
action_tokenizer: The internal tokenizer/processor instance, loaded during initialization.
|
||||
paligemma_tokenizer_name: The name of a pretrained PaliGemma tokenizer from the Hugging Face Hub (e.g., "google/paligemma-3b-pt-224").
|
||||
**Attributes**:
|
||||
- **tokenizer_name** -- The name of a pretrained processor from the Hugging Face Hub (e.g.,
|
||||
"lerobot/fast-action-tokenizer").
|
||||
- **tokenizer** -- A pre-initialized processor/tokenizer object. If provided, `tokenizer_name` is
|
||||
ignored.
|
||||
- **trust_remote_code** (`bool`) -- Whether to trust remote code when loading the tokenizer (required
|
||||
for some tokenizers).
|
||||
- **action_tokenizer** (`Any`) -- The internal tokenizer/processor instance, loaded during
|
||||
initialization.
|
||||
- **paligemma_tokenizer_name** (`str`) -- The name of a pretrained PaliGemma tokenizer from the
|
||||
Hugging Face Hub (e.g., "google/paligemma-3b-pt-224").
|
||||
"""
|
||||
|
||||
action_tokenizer_name: str | None = None
|
||||
|
||||
+59
-23
@@ -13,8 +13,7 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
Actor server runner for distributed HILSerl robot policy training.
|
||||
"""Actor server runner for distributed HILSerl robot policy training.
|
||||
|
||||
This script implements the actor component of the distributed HILSerl architecture.
|
||||
It executes the policy in the robot environment, collects experience,
|
||||
@@ -119,6 +118,16 @@ from .train_rl import TrainRLServerPipelineConfig
|
||||
|
||||
@parser.wrap()
|
||||
def actor_cli(cfg: TrainRLServerPipelineConfig):
|
||||
"""CLI entry point for the HILSerl actor server.
|
||||
|
||||
Connects to the learner server over gRPC, then launches (as threads or processes, depending on
|
||||
`cfg.policy.concurrency.multiprocessing_context`) the background workers that receive updated
|
||||
policy parameters and stream transitions/interactions back to the learner, while running the
|
||||
policy-environment interaction loop (`act_with_policy`) on the main thread/process.
|
||||
|
||||
Args:
|
||||
cfg (`TrainRLServerPipelineConfig`): Parsed from the CLI.
|
||||
"""
|
||||
# Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
|
||||
require_package("grpcio", extra="hilserl", import_name="grpc")
|
||||
cfg.validate()
|
||||
@@ -234,18 +243,19 @@ def act_with_policy(
|
||||
transitions_queue: Queue,
|
||||
interactions_queue: Queue,
|
||||
):
|
||||
"""
|
||||
Executes policy interaction within the environment.
|
||||
"""Executes policy interaction within the environment.
|
||||
|
||||
This function rolls out the policy in the environment, collecting interaction data and pushing it to a queue for streaming to the learner.
|
||||
Once an episode is completed, updated network parameters received from the learner are retrieved from a queue and loaded into the network.
|
||||
|
||||
Args:
|
||||
cfg: Configuration settings for the interaction process.
|
||||
shutdown_event: Event to check if the process should shutdown.
|
||||
parameters_queue: Queue to receive updated network parameters from the learner.
|
||||
transitions_queue: Queue to send transitions to the learner.
|
||||
interactions_queue: Queue to send interactions to the learner.
|
||||
cfg (`TrainRLServerPipelineConfig`): Training configuration.
|
||||
shutdown_event (`Event`): Set to stop the policy loop.
|
||||
parameters_queue (`Queue`): Queue of serialized learner weights, drained via
|
||||
`update_policy_parameters`.
|
||||
transitions_queue (`Queue`): Queue transitions are pushed to for streaming to the learner.
|
||||
interactions_queue (`Queue`): Queue interaction messages are pushed to for streaming to the
|
||||
learner.
|
||||
"""
|
||||
# Initialize logging for multiprocessing
|
||||
if not use_threads(cfg):
|
||||
@@ -440,7 +450,8 @@ def establish_learner_connection(
|
||||
Args:
|
||||
stub (services_pb2_grpc.LearnerServiceStub): The stub to use for the connection.
|
||||
shutdown_event (Event): The event to check if the connection should be established.
|
||||
attempts (int): The number of attempts to establish the connection.
|
||||
attempts (int, *optional*, defaults to 30): The number of attempts to establish the connection.
|
||||
|
||||
Returns:
|
||||
bool: True if the connection is established, False otherwise.
|
||||
"""
|
||||
@@ -473,7 +484,6 @@ def learner_service_client(
|
||||
Returns:
|
||||
tuple[services_pb2_grpc.LearnerServiceStub, grpc.Channel]: The stub and the channel.
|
||||
"""
|
||||
|
||||
channel = grpc.insecure_channel(
|
||||
f"{host}:{port}",
|
||||
grpc_channel_options(),
|
||||
@@ -496,8 +506,8 @@ def receive_policy(
|
||||
cfg (TrainRLServerPipelineConfig): The configuration for the actor.
|
||||
parameters_queue (Queue): The queue to receive the parameters.
|
||||
shutdown_event (Event): The event to check if the process should shutdown.
|
||||
learner_client (services_pb2_grpc.LearnerServiceStub | None): Optional pre-created stub.
|
||||
grpc_channel (grpc.Channel | None): Optional pre-created channel.
|
||||
learner_client (services_pb2_grpc.LearnerServiceStub | None, *optional*): Optional pre-created stub.
|
||||
grpc_channel (grpc.Channel | None, *optional*): Optional pre-created channel.
|
||||
"""
|
||||
logging.info("[ACTOR] Start receiving parameters from the Learner")
|
||||
if not use_threads(cfg):
|
||||
@@ -557,10 +567,9 @@ def send_transitions(
|
||||
cfg (TrainRLServerPipelineConfig): The configuration for the actor.
|
||||
transitions_queue (Queue): The queue to receive the transitions.
|
||||
shutdown_event (Event): The event to check if the process should shutdown.
|
||||
learner_client (services_pb2_grpc.LearnerServiceStub | None): Optional pre-created stub.
|
||||
grpc_channel (grpc.Channel | None): Optional pre-created channel.
|
||||
learner_client (services_pb2_grpc.LearnerServiceStub | None, *optional*): Optional pre-created stub.
|
||||
grpc_channel (grpc.Channel | None, *optional*): Optional pre-created channel.
|
||||
"""
|
||||
|
||||
if not use_threads(cfg):
|
||||
# Create a process-specific log file
|
||||
log_dir = os.path.join(cfg.output_dir, "logs")
|
||||
@@ -612,10 +621,9 @@ def send_interactions(
|
||||
cfg (TrainRLServerPipelineConfig): The configuration for the actor.
|
||||
interactions_queue (Queue): The queue to receive the interactions.
|
||||
shutdown_event (Event): The event to check if the process should shutdown.
|
||||
learner_client (services_pb2_grpc.LearnerServiceStub | None): Optional pre-created stub.
|
||||
grpc_channel (grpc.Channel | None): Optional pre-created channel.
|
||||
learner_client (services_pb2_grpc.LearnerServiceStub | None, *optional*): Optional pre-created stub.
|
||||
grpc_channel (grpc.Channel | None, *optional*): Optional pre-created channel.
|
||||
"""
|
||||
|
||||
if not use_threads(cfg):
|
||||
# Create a process-specific log file
|
||||
log_dir = os.path.join(cfg.output_dir, "logs")
|
||||
@@ -657,6 +665,17 @@ def transitions_stream(
|
||||
transitions_queue: Queue,
|
||||
timeout: float,
|
||||
) -> "Generator[Any, None, services_pb2.Empty]":
|
||||
"""GRPC client-streaming generator that forwards queued transitions to the learner.
|
||||
|
||||
Args:
|
||||
shutdown_event (`Event`): Set to stop streaming and return.
|
||||
transitions_queue (`Queue`): Queue of serialized transition batches, filled by
|
||||
`push_transitions_to_transport_queue`.
|
||||
timeout (`float`): Seconds to wait for a queue item before checking `shutdown_event` again.
|
||||
|
||||
Yields:
|
||||
Chunks of a `services_pb2.Transition` message, produced by `send_bytes_in_chunks`.
|
||||
"""
|
||||
while not shutdown_event.is_set():
|
||||
try:
|
||||
message = transitions_queue.get(block=True, timeout=timeout)
|
||||
@@ -676,6 +695,16 @@ def interactions_stream(
|
||||
interactions_queue: Queue,
|
||||
timeout: float,
|
||||
) -> "Generator[Any, None, services_pb2.Empty]":
|
||||
"""GRPC client-streaming generator that forwards queued interaction messages to the learner.
|
||||
|
||||
Args:
|
||||
shutdown_event (`Event`): Set to stop streaming and return.
|
||||
interactions_queue (`Queue`): Queue of serialized interaction messages.
|
||||
timeout (`float`): Seconds to wait for a queue item before checking `shutdown_event` again.
|
||||
|
||||
Yields:
|
||||
Chunks of a `services_pb2.InteractionMessage`, produced by `send_bytes_in_chunks`.
|
||||
"""
|
||||
while not shutdown_event.is_set():
|
||||
try:
|
||||
message = interactions_queue.get(block=True, timeout=timeout)
|
||||
@@ -718,12 +747,11 @@ def update_policy_parameters(algorithm: RLAlgorithm, parameters_queue: Queue, de
|
||||
|
||||
|
||||
def push_transitions_to_transport_queue(transitions: list, transitions_queue):
|
||||
"""Send transitions to learner in smaller chunks to avoid network issues.
|
||||
"""Move `transitions` to CPU, check for NaNs, and enqueue them for the learner.
|
||||
|
||||
Args:
|
||||
transitions: List of transitions to send
|
||||
message_queue: Queue to send messages to learner
|
||||
chunk_size: Size of each chunk to send
|
||||
transitions (`list`): Transitions to send, as produced by the actor's rollout loop.
|
||||
transitions_queue (`Queue`): Queue drained by `transitions_stream`.
|
||||
"""
|
||||
transition_to_send_to_learner = []
|
||||
for transition in transitions:
|
||||
@@ -760,6 +788,13 @@ def get_frequency_stats(timer: TimerManager) -> dict[str, float]:
|
||||
|
||||
|
||||
def log_policy_frequency_issue(policy_fps: float, cfg: TrainRLServerPipelineConfig, interaction_step: int):
|
||||
"""Log a warning if `policy_fps` is below the environment's target `cfg.env.fps`.
|
||||
|
||||
Args:
|
||||
policy_fps (`float`): Measured policy loop frequency.
|
||||
cfg (`TrainRLServerPipelineConfig`): Provides the target `cfg.env.fps` to compare against.
|
||||
interaction_step (`int`): Current interaction step, included in the warning message.
|
||||
"""
|
||||
if policy_fps < cfg.env.fps:
|
||||
logging.warning(
|
||||
f"[ACTOR] Policy FPS {policy_fps:.1f} below required {cfg.env.fps} at step {interaction_step}"
|
||||
@@ -767,6 +802,7 @@ def log_policy_frequency_issue(policy_fps: float, cfg: TrainRLServerPipelineConf
|
||||
|
||||
|
||||
def use_threads(cfg: TrainRLServerPipelineConfig) -> bool:
|
||||
"""Whether the actor's background workers should run as threads instead of processes."""
|
||||
return cfg.policy.concurrency.actor == "threads"
|
||||
|
||||
|
||||
|
||||
@@ -101,6 +101,7 @@ class RLAlgorithm(HubMixin, abc.ABC):
|
||||
|
||||
@optimization_step.setter
|
||||
def optimization_step(self, value: int) -> None:
|
||||
"""Set the current learner optimization step."""
|
||||
self._optimization_step = int(value)
|
||||
|
||||
def get_weights(self) -> dict[str, Any]:
|
||||
|
||||
@@ -45,7 +45,6 @@ class TrainingStats:
|
||||
|
||||
def to_log_dict(self) -> dict[str, float]:
|
||||
"""Flatten all stats into a single dict for logging."""
|
||||
|
||||
d: dict[str, float] = {}
|
||||
for name, val in self.losses.items():
|
||||
d[name] = val
|
||||
@@ -98,6 +97,35 @@ class RLAlgorithmConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
|
||||
revision: str | None = None,
|
||||
**algo_kwargs: Any,
|
||||
) -> T:
|
||||
"""Load an algorithm config from a local directory or the Hugging Face Hub.
|
||||
|
||||
Args:
|
||||
pretrained_name_or_path (`str | Path`):
|
||||
Local directory containing `config.json`, or a Hub repo id.
|
||||
force_download (`bool`, *optional*, defaults to `False`):
|
||||
Whether to force re-download the config even if it's cached.
|
||||
resume_download (`bool | None`, *optional*):
|
||||
Whether to resume an interrupted download.
|
||||
proxies (`dict[Any, Any] | None`, *optional*):
|
||||
Proxies to use for the download request.
|
||||
token (`str | bool | None`, *optional*):
|
||||
Hugging Face Hub authentication token.
|
||||
cache_dir (`str | Path | None`, *optional*):
|
||||
Directory to cache the downloaded config in.
|
||||
local_files_only (`bool`, *optional*, defaults to `False`):
|
||||
Whether to only look for files locally, without querying the Hub.
|
||||
revision (`str | None`, *optional*):
|
||||
Hub revision (branch, tag, or commit hash) to load from.
|
||||
**algo_kwargs: Attribute overrides applied to the loaded config instance.
|
||||
|
||||
Returns:
|
||||
RLAlgorithmConfig: The loaded config, as the concrete registered subclass.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If no `config.json` is found locally or on the Hub.
|
||||
TypeError: If loaded via a specific subclass but the config's registered type doesn't
|
||||
match it.
|
||||
"""
|
||||
model_id = str(pretrained_name_or_path)
|
||||
config_file: str | None = None
|
||||
if Path(model_id).is_dir():
|
||||
|
||||
@@ -24,8 +24,8 @@ def make_algorithm_config(algorithm_type: str, **kwargs) -> RLAlgorithmConfig:
|
||||
"""Instantiate an `RLAlgorithmConfig` from its registered type name.
|
||||
|
||||
Args:
|
||||
algorithm_type: Registry key of the algorithm (e.g. ``"sac"``).
|
||||
**kwargs: Keyword arguments forwarded to the config class constructor.
|
||||
algorithm_type (`str`): Registry key of the algorithm (e.g. `"sac"`).
|
||||
kwargs (`Any`, *optional*): Keyword arguments forwarded to the config class constructor.
|
||||
|
||||
Returns:
|
||||
An instance of the matching ``RLAlgorithmConfig`` subclass.
|
||||
@@ -44,14 +44,13 @@ def make_algorithm_config(algorithm_type: str, **kwargs) -> RLAlgorithmConfig:
|
||||
|
||||
|
||||
def get_algorithm_class(name: str) -> type[RLAlgorithm]:
|
||||
"""
|
||||
Retrieves an RL algorithm class by its registered name.
|
||||
"""Retrieves an RL algorithm class by its registered name.
|
||||
|
||||
This function uses dynamic imports to avoid loading all algorithm classes into
|
||||
memory at once, improving startup time and reducing dependencies.
|
||||
|
||||
Args:
|
||||
name: The name of the algorithm. Supported names are "sac".
|
||||
name (`str`): The name of the algorithm. Supported names are "sac".
|
||||
|
||||
Returns:
|
||||
The algorithm class corresponding to the given name.
|
||||
@@ -70,8 +69,7 @@ def get_algorithm_class(name: str) -> type[RLAlgorithm]:
|
||||
|
||||
|
||||
def make_algorithm(cfg: RLAlgorithmConfig, policy: torch.nn.Module) -> RLAlgorithm:
|
||||
"""
|
||||
Instantiate an RL algorithm.
|
||||
"""Instantiate an RL algorithm.
|
||||
|
||||
This factory function looks up the :class:`RLAlgorithm` subclass that matches
|
||||
``cfg.type`` and instantiates it with the provided policy. It also enforces
|
||||
@@ -79,8 +77,8 @@ def make_algorithm(cfg: RLAlgorithmConfig, policy: torch.nn.Module) -> RLAlgorit
|
||||
normally handled by :meth:`TrainRLServerPipelineConfig.validate`).
|
||||
|
||||
Args:
|
||||
cfg: The algorithm configuration. Must have ``policy_config`` set.
|
||||
policy: The policy module the algorithm will train.
|
||||
cfg (`RLAlgorithmConfig`): The algorithm configuration. Must have `policy_config` set.
|
||||
policy (`torch.nn.Module`): The policy module the algorithm will train.
|
||||
|
||||
Returns:
|
||||
An instantiated :class:`RLAlgorithm`.
|
||||
|
||||
@@ -39,52 +39,73 @@ class SACAlgorithmConfig(RLAlgorithmConfig):
|
||||
update loop. The policy-side (actor + observation encoder) lives in
|
||||
:class:`~lerobot.policies.gaussian_actor.GaussianActorConfig` and is
|
||||
referenced via :attr:`policy_config`.
|
||||
|
||||
Args:
|
||||
actor_lr (`float`, *optional*, defaults to 0.0003):
|
||||
Learning rate for the actor network.
|
||||
critic_lr (`float`, *optional*, defaults to 0.0003):
|
||||
Learning rate for the critic network.
|
||||
temperature_lr (`float`, *optional*, defaults to 0.0003):
|
||||
Learning rate for the temperature parameter.
|
||||
discount (`float`, *optional*, defaults to 0.99):
|
||||
Discount factor for the Bellman update.
|
||||
use_backup_entropy (`bool`, *optional*, defaults to `True`):
|
||||
Whether to use backup entropy in the Bellman target.
|
||||
critic_target_update_weight (`float`, *optional*, defaults to 0.005):
|
||||
Polyak-averaging weight for the critic target update.
|
||||
num_critics (`int`, *optional*, defaults to 2):
|
||||
Number of critics in the ensemble.
|
||||
num_subsample_critics (`int | None`, *optional*):
|
||||
Number of critics to subsample from the ensemble for each Bellman target computation.
|
||||
`None` uses the full ensemble.
|
||||
critic_network_kwargs (`CriticNetworkConfig`, *optional*):
|
||||
Configuration for the (continuous-action) critic network architecture.
|
||||
discrete_critic_network_kwargs (`CriticNetworkConfig`, *optional*):
|
||||
Configuration for the discrete-action critic network architecture.
|
||||
temperature_init (`float`, *optional*, defaults to 1.0):
|
||||
Initial value of the entropy temperature.
|
||||
target_entropy (`float | None`, *optional*):
|
||||
Target entropy for automatic temperature tuning. If `None`, defaults to `-|A|/2` where
|
||||
`|A|` is the total action dimension (continuous + 1 if there is a discrete action head).
|
||||
utd_ratio (`int`, *optional*, defaults to 1):
|
||||
Update-to-data ratio. Set to `>1` to enable extra critic updates per env step.
|
||||
policy_update_freq (`int`, *optional*, defaults to 1):
|
||||
Frequency of policy updates, in units of critic updates.
|
||||
grad_clip_norm (`float`, *optional*, defaults to 40.0):
|
||||
Gradient-clipping norm applied during optimization.
|
||||
use_torch_compile (`bool`, *optional*, defaults to `False`):
|
||||
Whether to `torch.compile` the algorithm's forward passes. Currently disabled by default.
|
||||
policy_config (`PreTrainedConfig | None`, *optional*):
|
||||
The policy (actor) config this algorithm trains. Populated via `from_policy_config` or by
|
||||
`TrainRLServerPipelineConfig.validate` before the algorithm is constructed.
|
||||
"""
|
||||
|
||||
# Optimizer learning rates
|
||||
# Learning rate for the actor network
|
||||
actor_lr: float = 3e-4
|
||||
# Learning rate for the critic network
|
||||
critic_lr: float = 3e-4
|
||||
# Learning rate for the temperature parameter
|
||||
temperature_lr: float = 3e-4
|
||||
|
||||
# Bellman update
|
||||
# Discount factor for the SAC algorithm
|
||||
discount: float = 0.99
|
||||
# Whether to use backup entropy for the SAC algorithm
|
||||
use_backup_entropy: bool = True
|
||||
# Weight for the critic target update
|
||||
critic_target_update_weight: float = 0.005
|
||||
|
||||
# Critic ensemble
|
||||
# Number of critics in the ensemble
|
||||
num_critics: int = 2
|
||||
# Number of subsampled critics for training
|
||||
num_subsample_critics: int | None = None
|
||||
# Configuration for the critic network architecture
|
||||
critic_network_kwargs: CriticNetworkConfig = field(default_factory=CriticNetworkConfig)
|
||||
# Configuration for the discrete critic network
|
||||
discrete_critic_network_kwargs: CriticNetworkConfig = field(default_factory=CriticNetworkConfig)
|
||||
|
||||
# Temperature / entropy
|
||||
# Initial temperature value
|
||||
temperature_init: float = 1.0
|
||||
# Target entropy for automatic temperature tuning. If ``None``, defaults to
|
||||
# ``-|A|/2`` where ``|A|`` is the total action dimension (continuous + 1 if
|
||||
# there is a discrete action head).
|
||||
target_entropy: float | None = None
|
||||
|
||||
# Update loop
|
||||
# Update-to-data ratio. Set to >1 to enable extra critic updates per env step.
|
||||
utd_ratio: int = 1
|
||||
# Frequency of policy updates
|
||||
policy_update_freq: int = 1
|
||||
# Gradient clipping norm for the SAC algorithm
|
||||
grad_clip_norm: float = 40.0
|
||||
|
||||
# Optimizations
|
||||
# torch.compile is currently disabled by default
|
||||
use_torch_compile: bool = False
|
||||
|
||||
# Policy config
|
||||
|
||||
@@ -55,6 +55,15 @@ class SACAlgorithm(RLAlgorithm):
|
||||
policy: GaussianActorPolicy,
|
||||
config: SACAlgorithmConfig,
|
||||
):
|
||||
"""Build the critic ensemble, target networks, and temperature from `config`.
|
||||
|
||||
Args:
|
||||
policy (`GaussianActorPolicy`):
|
||||
The actor policy this algorithm trains. Its observation encoder is shared with the
|
||||
critics.
|
||||
config (`SACAlgorithmConfig`):
|
||||
Algorithm configuration.
|
||||
"""
|
||||
self.config = config
|
||||
self.policy_config = config.policy_config
|
||||
self.policy = policy
|
||||
@@ -144,17 +153,18 @@ class SACAlgorithm(RLAlgorithm):
|
||||
use_target: bool = False,
|
||||
observation_features: Tensor | None = None,
|
||||
) -> Tensor:
|
||||
"""Forward pass through a critic network ensemble
|
||||
"""Forward pass through a critic network ensemble.
|
||||
|
||||
Args:
|
||||
observations: Dictionary of observations
|
||||
actions: Action tensor
|
||||
use_target: If True, use target critics, otherwise use ensemble critics
|
||||
observation_features: Optional pre-computed observation features to avoid recomputing
|
||||
encoder output
|
||||
|
||||
Returns:
|
||||
Tensor of Q-values from all critics
|
||||
"""
|
||||
|
||||
critics = self.critic_target if use_target else self.critic_ensemble
|
||||
q_values = critics(observations, actions, observation_features)
|
||||
return q_values
|
||||
@@ -162,7 +172,7 @@ class SACAlgorithm(RLAlgorithm):
|
||||
def _discrete_critic_forward(
|
||||
self, observations, use_target=False, observation_features=None
|
||||
) -> torch.Tensor:
|
||||
"""Forward pass through a discrete critic network
|
||||
"""Forward pass through a discrete critic network.
|
||||
|
||||
Args:
|
||||
observations: Dictionary of observations
|
||||
@@ -408,7 +418,7 @@ class SACAlgorithm(RLAlgorithm):
|
||||
return actor_loss
|
||||
|
||||
def _compute_loss_temperature(self, batch: dict[str, Any]) -> Tensor:
|
||||
"""Compute the temperature loss"""
|
||||
"""Compute the temperature loss."""
|
||||
observations = batch["state"]
|
||||
observation_features = batch.get("observation_feature")
|
||||
|
||||
@@ -420,7 +430,7 @@ class SACAlgorithm(RLAlgorithm):
|
||||
return temperature_loss
|
||||
|
||||
def _update_target_networks(self) -> None:
|
||||
"""Update target networks with exponential moving average"""
|
||||
"""Update target networks with exponential moving average."""
|
||||
for target_p, p in zip(
|
||||
self.critic_target.parameters(), self.critic_ensemble.parameters(), strict=True
|
||||
):
|
||||
@@ -461,8 +471,7 @@ class SACAlgorithm(RLAlgorithm):
|
||||
return forward_batch
|
||||
|
||||
def make_optimizers_and_scheduler(self) -> dict[str, Optimizer]:
|
||||
"""
|
||||
Creates and returns optimizers for the actor, critic, and temperature components of a reinforcement learning policy.
|
||||
"""Creates and returns optimizers for the actor, critic, and temperature components of a reinforcement learning policy.
|
||||
|
||||
This function sets up Adam optimizers for:
|
||||
- The **actor network**, ensuring that only relevant parameters are optimized.
|
||||
@@ -471,7 +480,7 @@ class SACAlgorithm(RLAlgorithm):
|
||||
|
||||
It also initializes a learning rate scheduler, though currently, it is set to `None`.
|
||||
|
||||
NOTE:
|
||||
Note:
|
||||
- If the encoder is shared, its parameters are excluded from the actor's optimization process.
|
||||
- The policy's log temperature (`log_alpha`) is wrapped in a list to ensure proper optimization as a standalone tensor.
|
||||
|
||||
@@ -496,6 +505,7 @@ class SACAlgorithm(RLAlgorithm):
|
||||
return self.optimizers
|
||||
|
||||
def get_optimizers(self) -> dict[str, Optimizer]:
|
||||
"""See [`~rl.algorithms.RLAlgorithm.get_optimizers`]."""
|
||||
return self.optimizers
|
||||
|
||||
def get_weights(self) -> dict[str, Any]:
|
||||
@@ -560,20 +570,18 @@ class SACAlgorithm(RLAlgorithm):
|
||||
def get_observation_features(
|
||||
self, observations: Tensor, next_observations: Tensor
|
||||
) -> tuple[Tensor | None, Tensor | None]:
|
||||
"""
|
||||
Get observation features from the policy encoder. It act as cache for the observation features.
|
||||
when the encoder is frozen, the observation features are not updated.
|
||||
We can save compute by caching the observation features.
|
||||
"""Get observation features from the policy encoder, acting as a cache.
|
||||
|
||||
When the encoder is frozen, the observation features are not updated, so we can save compute
|
||||
by caching them here instead of recomputing on every critic/actor forward pass.
|
||||
|
||||
Args:
|
||||
policy: The policy model
|
||||
observations: The current observations
|
||||
next_observations: The next observations
|
||||
|
||||
Returns:
|
||||
tuple: observation_features, next_observation_features
|
||||
"""
|
||||
|
||||
if self.policy.config.vision_encoder_name is None or not self.policy.config.freeze_vision_encoder:
|
||||
return None, None
|
||||
|
||||
@@ -595,6 +603,8 @@ def _split_prefix(state: dict[str, torch.Tensor], prefix: str) -> dict[str, torc
|
||||
|
||||
|
||||
class CriticHead(nn.Module):
|
||||
"""A single Q-value head: an MLP followed by a scalar linear output layer."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dim: int,
|
||||
@@ -605,6 +615,23 @@ class CriticHead(nn.Module):
|
||||
init_final: float | None = None,
|
||||
final_activation: Callable[[torch.Tensor], torch.Tensor] | str | None = None,
|
||||
):
|
||||
"""Build the MLP trunk and scalar output layer.
|
||||
|
||||
Args:
|
||||
input_dim (`int`): Dimension of the concatenated observation-encoding + action input.
|
||||
hidden_dims (`list[int]`): Hidden layer widths of the MLP trunk.
|
||||
activations (`Callable[[torch.Tensor], torch.Tensor] | str`, *optional*, defaults to `nn.SiLU()`):
|
||||
Activation used between hidden layers.
|
||||
activate_final (`bool`, *optional*, defaults to `False`): Whether to apply `activations`
|
||||
after the last hidden layer.
|
||||
dropout_rate (`float | None`, *optional*): Dropout probability applied between hidden
|
||||
layers. `None` disables dropout.
|
||||
init_final (`float | None`, *optional*): When set, the output layer's weight and bias are
|
||||
initialized uniformly in `[-init_final, init_final]` instead of the default
|
||||
orthogonal initialization.
|
||||
final_activation (`Callable[[torch.Tensor], torch.Tensor] | str | None`, *optional*):
|
||||
Activation applied after the MLP trunk's last hidden layer, before the output layer.
|
||||
"""
|
||||
super().__init__()
|
||||
self.net = MLP(
|
||||
input_dim=input_dim,
|
||||
@@ -622,17 +649,17 @@ class CriticHead(nn.Module):
|
||||
orthogonal_init()(self.output_layer.weight)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Compute the scalar Q-value for `x` (a concatenated observation-encoding + action tensor)."""
|
||||
return self.output_layer(self.net(x))
|
||||
|
||||
|
||||
class CriticEnsemble(nn.Module):
|
||||
"""
|
||||
CriticEnsemble wraps multiple CriticHead modules into an ensemble.
|
||||
"""CriticEnsemble wraps multiple CriticHead modules into an ensemble.
|
||||
|
||||
Args:
|
||||
encoder (GaussianActorObservationEncoder): encoder for observations.
|
||||
ensemble (List[CriticHead]): list of critic heads.
|
||||
init_final (float | None): optional initializer scale for final layers.
|
||||
init_final (float | None, *optional*): optional initializer scale for final layers.
|
||||
|
||||
Forward returns a tensor of shape (num_critics, batch_size) containing Q-values.
|
||||
"""
|
||||
@@ -643,6 +670,14 @@ class CriticEnsemble(nn.Module):
|
||||
ensemble: list[CriticHead],
|
||||
init_final: float | None = None,
|
||||
):
|
||||
"""Wrap `ensemble` behind the shared `encoder`.
|
||||
|
||||
Args:
|
||||
encoder (`GaussianActorObservationEncoder`): Shared observation encoder for all critics.
|
||||
ensemble (`list[CriticHead]`): The critic heads making up the ensemble.
|
||||
init_final (`float | None`, *optional*): Stored for introspection; each `CriticHead` is
|
||||
already initialized with it before being passed in here.
|
||||
"""
|
||||
super().__init__()
|
||||
self.encoder = encoder
|
||||
self.init_final = init_final
|
||||
@@ -654,6 +689,19 @@ class CriticEnsemble(nn.Module):
|
||||
actions: torch.Tensor,
|
||||
observation_features: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Encode `observations` and return each ensemble member's Q-value for `actions`.
|
||||
|
||||
Args:
|
||||
observations (`dict[str, torch.Tensor]`): Raw observation tensors, moved to the module's
|
||||
device.
|
||||
actions (`torch.Tensor`): Action tensor to evaluate.
|
||||
observation_features (`torch.Tensor | None`, *optional*): Pre-computed encoder output,
|
||||
e.g. from `SACAlgorithm.get_observation_features`. Bypasses re-encoding when the
|
||||
vision encoder is frozen.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Q-values of shape `(num_critics, batch_size)`.
|
||||
"""
|
||||
device = get_device_from_parameters(self)
|
||||
# Move each tensor in observations to device
|
||||
observations = {k: v.to(device) for k, v in observations.items()}
|
||||
|
||||
+34
-27
@@ -30,6 +30,19 @@ from lerobot.utils.transition import Transition
|
||||
|
||||
|
||||
class BatchTransition(TypedDict):
|
||||
"""A batch of transitions sampled from a `ReplayBuffer`.
|
||||
|
||||
**Attributes**:
|
||||
- **state** (`dict[str, torch.Tensor]`) -- Batched observation tensors at time `t`.
|
||||
- **action** (`torch.Tensor`) -- Batched actions taken at time `t`.
|
||||
- **reward** (`torch.Tensor`) -- Batched rewards received after `action`.
|
||||
- **next_state** (`dict[str, torch.Tensor]`) -- Batched observation tensors at time `t+1`.
|
||||
- **done** (`torch.Tensor`) -- Batched episode-termination flags.
|
||||
- **truncated** (`torch.Tensor`) -- Batched episode-truncation flags.
|
||||
- **complementary_info** (`dict[str, torch.Tensor | float | int] | None`) -- Optional extra
|
||||
per-transition data (e.g. intervention flags), when present in the underlying dataset.
|
||||
"""
|
||||
|
||||
state: dict[str, torch.Tensor]
|
||||
action: torch.Tensor
|
||||
reward: torch.Tensor
|
||||
@@ -40,10 +53,7 @@ class BatchTransition(TypedDict):
|
||||
|
||||
|
||||
def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
|
||||
"""
|
||||
Perform a per-image random crop over a batch of images in a vectorized way.
|
||||
(Same as shown previously.)
|
||||
"""
|
||||
"""Perform a per-image random crop over a batch of images in a vectorized way."""
|
||||
B, C, H, W = images.shape # noqa: N806
|
||||
crop_h, crop_w = output_size
|
||||
|
||||
@@ -72,13 +82,15 @@ def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Te
|
||||
|
||||
|
||||
def random_shift(images: torch.Tensor, pad: int = 4):
|
||||
"""Vectorized random shift, imgs: (B,C,H,W), pad: #pixels"""
|
||||
"""Vectorized random shift. `images` has shape `(B, C, H, W)`; `pad` is the shift range in pixels."""
|
||||
_, _, h, w = images.shape
|
||||
images = F.pad(input=images, pad=(pad, pad, pad, pad), mode="replicate")
|
||||
return random_crop_vectorized(images=images, output_size=(h, w))
|
||||
|
||||
|
||||
class ReplayBuffer:
|
||||
"""In-memory replay buffer of `Transition`s, sampled in batches for off-policy RL training."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
capacity: int,
|
||||
@@ -89,11 +101,12 @@ class ReplayBuffer:
|
||||
storage_device: str = "cpu",
|
||||
optimize_memory: bool = False,
|
||||
):
|
||||
"""
|
||||
Replay buffer for storing transitions.
|
||||
"""Replay buffer for storing transitions.
|
||||
|
||||
It will allocate tensors on the specified device, when the first transition is added.
|
||||
NOTE: If you encounter memory issues, you can try to use the `optimize_memory` flag to save memory or
|
||||
and use the `storage_device` flag to store the buffer on a different device.
|
||||
|
||||
Args:
|
||||
capacity (int): Maximum number of transitions to store in the buffer.
|
||||
device (str): The device where the tensors will be moved when sampling ("cuda:0" or "cpu").
|
||||
@@ -187,6 +200,7 @@ class ReplayBuffer:
|
||||
self.initialized = True
|
||||
|
||||
def __len__(self):
|
||||
"""Number of transitions currently stored in the buffer."""
|
||||
return self.size
|
||||
|
||||
def add(
|
||||
@@ -305,8 +319,8 @@ class ReplayBuffer:
|
||||
async_prefetch: bool = True,
|
||||
queue_size: int = 2,
|
||||
):
|
||||
"""
|
||||
Creates an infinite iterator that yields batches of transitions.
|
||||
"""Creates an infinite iterator that yields batches of transitions.
|
||||
|
||||
Will automatically restart when internal iterator is exhausted.
|
||||
|
||||
Args:
|
||||
@@ -329,10 +343,9 @@ class ReplayBuffer:
|
||||
yield from iterator
|
||||
|
||||
def _get_async_iterator(self, batch_size: int, queue_size: int = 2):
|
||||
"""
|
||||
Create an iterator that continuously yields prefetched batches in a
|
||||
background thread. The design is intentionally simple and avoids busy
|
||||
waiting / complex state management.
|
||||
"""Create an iterator that continuously yields prefetched batches in a background thread.
|
||||
|
||||
The design is intentionally simple and avoids busy waiting / complex state management.
|
||||
|
||||
Args:
|
||||
batch_size (int): Size of batches to sample.
|
||||
@@ -383,8 +396,7 @@ class ReplayBuffer:
|
||||
producer_thread.join(timeout=1.0)
|
||||
|
||||
def _get_naive_iterator(self, batch_size: int, queue_size: int = 2):
|
||||
"""
|
||||
Creates a simple non-threaded iterator that yields batches.
|
||||
"""Creates a simple non-threaded iterator that yields batches.
|
||||
|
||||
Args:
|
||||
batch_size (int): Size of batches to sample
|
||||
@@ -398,6 +410,7 @@ class ReplayBuffer:
|
||||
queue = collections.deque()
|
||||
|
||||
def enqueue(n):
|
||||
"""Sample `n` more batches and append them to `queue`."""
|
||||
for _ in range(n):
|
||||
data = self.sample(batch_size)
|
||||
queue.append(data)
|
||||
@@ -419,8 +432,7 @@ class ReplayBuffer:
|
||||
storage_device: str = "cpu",
|
||||
optimize_memory: bool = False,
|
||||
) -> "ReplayBuffer":
|
||||
"""
|
||||
Convert a LeRobotDataset into a ReplayBuffer.
|
||||
"""Convert a LeRobotDataset into a ReplayBuffer.
|
||||
|
||||
Args:
|
||||
lerobot_dataset (LeRobotDataset): The dataset to convert.
|
||||
@@ -509,9 +521,7 @@ class ReplayBuffer:
|
||||
root=None,
|
||||
task_name="from_replay_buffer",
|
||||
) -> LeRobotDataset:
|
||||
"""
|
||||
Converts all transitions in this ReplayBuffer into a single LeRobotDataset object.
|
||||
"""
|
||||
"""Converts all transitions in this ReplayBuffer into a single LeRobotDataset object."""
|
||||
if self.size == 0:
|
||||
raise ValueError("The replay buffer is empty. Cannot convert to a dataset.")
|
||||
|
||||
@@ -612,8 +622,7 @@ class ReplayBuffer:
|
||||
dataset: LeRobotDataset,
|
||||
state_keys: Sequence[str] | None = None,
|
||||
) -> list[Transition]:
|
||||
"""
|
||||
Convert a LeRobotDataset into a list of RL (s, a, r, s', done) transitions.
|
||||
"""Convert a LeRobotDataset into a list of RL (s, a, r, s', done) transitions.
|
||||
|
||||
Args:
|
||||
dataset (LeRobotDataset):
|
||||
@@ -733,12 +742,11 @@ class ReplayBuffer:
|
||||
|
||||
# Utility function to guess shapes/dtypes from a tensor
|
||||
def guess_feature_info(t, name: str):
|
||||
"""
|
||||
Return a dictionary with the 'dtype' and 'shape' for a given tensor or scalar value.
|
||||
"""Return a dictionary with the 'dtype' and 'shape' for a given tensor or scalar value.
|
||||
|
||||
If it looks like a 3D (C,H,W) shape, we might consider it an 'image'.
|
||||
Otherwise default to appropriate dtype for numeric.
|
||||
"""
|
||||
|
||||
shape = tuple(t.shape)
|
||||
# Basic guess: if we have exactly 3 dims and shape[0] in {1, 3}, guess 'image'
|
||||
if len(shape) == 3 and shape[0] in [1, 3]:
|
||||
@@ -757,8 +765,7 @@ def guess_feature_info(t, name: str):
|
||||
def concatenate_batch_transitions(
|
||||
left_batch_transitions: BatchTransition, right_batch_transition: BatchTransition
|
||||
) -> BatchTransition:
|
||||
"""
|
||||
Concatenates two BatchTransition objects into one.
|
||||
"""Concatenates two BatchTransition objects into one.
|
||||
|
||||
This function merges the right BatchTransition into the left one by concatenating
|
||||
all corresponding tensors along dimension 0. The operation modifies the left_batch_transitions
|
||||
|
||||
@@ -29,8 +29,7 @@ from lerobot.utils.constants import DONE, REWARD
|
||||
|
||||
|
||||
def select_rect_roi(img):
|
||||
"""
|
||||
Allows the user to draw a rectangular ROI on the image.
|
||||
"""Allows the user to draw a rectangular ROI on the image.
|
||||
|
||||
The user must click and drag to draw the rectangle.
|
||||
- While dragging, the rectangle is dynamically drawn.
|
||||
@@ -52,6 +51,7 @@ def select_rect_roi(img):
|
||||
index_x, index_y = -1, -1 # Initial click coordinates
|
||||
|
||||
def mouse_callback(event, x, y, flags, param):
|
||||
"""`cv2.setMouseCallback` handler that drives the click-and-drag ROI selection."""
|
||||
nonlocal index_x, index_y, drawing, roi, working_img
|
||||
|
||||
if event == cv2.EVENT_LBUTTONDOWN:
|
||||
@@ -118,12 +118,11 @@ def select_rect_roi(img):
|
||||
|
||||
|
||||
def select_square_roi_for_images(images: dict) -> dict:
|
||||
"""
|
||||
For each image in the provided dictionary, open a window to allow the user
|
||||
to select a rectangular ROI. Returns a dictionary mapping each key to a tuple
|
||||
(top, left, height, width) representing the ROI.
|
||||
"""For each image in the provided dictionary, open a window to allow the user to select a ROI.
|
||||
|
||||
Parameters:
|
||||
Returns a dictionary mapping each key to a tuple (top, left, height, width) representing the ROI.
|
||||
|
||||
Args:
|
||||
images (dict): Dictionary where keys are identifiers and values are OpenCV images.
|
||||
|
||||
Returns:
|
||||
@@ -149,9 +148,7 @@ def select_square_roi_for_images(images: dict) -> dict:
|
||||
|
||||
|
||||
def get_image_from_lerobot_dataset(dataset: LeRobotDataset):
|
||||
"""
|
||||
Find the first row in the dataset and extract the image in order to be used for the crop.
|
||||
"""
|
||||
"""Find the first row in the dataset and extract the image in order to be used for the crop."""
|
||||
row = dataset[0]
|
||||
image_dict = {}
|
||||
for k in row:
|
||||
@@ -169,19 +166,23 @@ def convert_lerobot_dataset_to_cropped_lerobot_dataset(
|
||||
push_to_hub: bool = False,
|
||||
task: str = "",
|
||||
) -> LeRobotDataset:
|
||||
"""
|
||||
Converts an existing LeRobotDataset by iterating over its episodes and frames,
|
||||
applying cropping and resizing to image observations, and saving a new dataset
|
||||
with the transformed data.
|
||||
"""Converts an existing LeRobotDataset to a new one with cropped/resized image observations.
|
||||
|
||||
Iterates over the source dataset's episodes and frames, applying cropping and resizing to image
|
||||
observations, and saves a new dataset with the transformed data.
|
||||
|
||||
Args:
|
||||
original_dataset (LeRobotDataset): The source dataset.
|
||||
crop_params_dict (dict[str, Tuple[int, int, int, int]]):
|
||||
original_dataset (`LeRobotDataset`): The source dataset.
|
||||
crop_params_dict (`dict[str, tuple[int, int, int, int]]`):
|
||||
A dictionary mapping observation keys to crop parameters (top, left, height, width).
|
||||
new_repo_id (str): Repository id for the new dataset.
|
||||
new_dataset_root (str): The root directory where the new dataset will be written.
|
||||
resize_size (tuple[int, int], optional): The target size (height, width) after cropping.
|
||||
Defaults to (128, 128).
|
||||
new_repo_id (`str`): Repository id for the new dataset.
|
||||
new_dataset_root (`str`): The root directory where the new dataset will be written.
|
||||
resize_size (`tuple[int, int]`, *optional*, defaults to `(128, 128)`): The target size
|
||||
(height, width) after cropping.
|
||||
push_to_hub (`bool`, *optional*, defaults to `False`): Whether to push the new dataset to the
|
||||
Hugging Face Hub.
|
||||
task (`str`, *optional*, defaults to `""`): Task description recorded on every frame of the
|
||||
new dataset.
|
||||
|
||||
Returns:
|
||||
LeRobotDataset: A new LeRobotDataset where the specified image observations have been cropped
|
||||
|
||||
@@ -49,6 +49,18 @@ class OnlineOfflineMixer(DataMixer):
|
||||
offline_buffer: ReplayBuffer | None = None,
|
||||
online_ratio: float = 1.0,
|
||||
):
|
||||
"""Create the mixer.
|
||||
|
||||
Args:
|
||||
online_buffer (`ReplayBuffer`): Buffer of transitions collected online during training.
|
||||
offline_buffer (`ReplayBuffer | None`, *optional*): Buffer of pre-collected offline
|
||||
transitions. When `None`, every batch is drawn from `online_buffer` alone.
|
||||
online_ratio (`float`, *optional*, defaults to 1.0): Fraction of each batch drawn from
|
||||
`online_buffer`; the remainder comes from `offline_buffer`. Must be in `[0, 1]`.
|
||||
|
||||
Raises:
|
||||
ValueError: If `online_ratio` is not in `[0, 1]`.
|
||||
"""
|
||||
if not 0.0 <= online_ratio <= 1.0:
|
||||
raise ValueError(f"online_ratio must be in [0, 1], got {online_ratio}")
|
||||
self.online_buffer = online_buffer
|
||||
@@ -56,6 +68,7 @@ class OnlineOfflineMixer(DataMixer):
|
||||
self.online_ratio = online_ratio
|
||||
|
||||
def sample(self, batch_size: int) -> BatchType:
|
||||
"""See [`~rl.data_sources.DataMixer.sample`]."""
|
||||
if self.offline_buffer is None:
|
||||
return self.online_buffer.sample(batch_size)
|
||||
|
||||
@@ -73,7 +86,6 @@ class OnlineOfflineMixer(DataMixer):
|
||||
queue_size: int = 2,
|
||||
):
|
||||
"""Yield batches by composing buffer async iterators."""
|
||||
|
||||
n_online = max(1, int(batch_size * self.online_ratio))
|
||||
|
||||
online_iter = self.online_buffer.get_iterator(
|
||||
|
||||
@@ -36,6 +36,13 @@ logging.basicConfig(level=logging.INFO)
|
||||
|
||||
|
||||
def eval_policy(env, policy, n_episodes):
|
||||
"""Roll out `policy` in `env` for `n_episodes` and log the per-episode and average reward.
|
||||
|
||||
Args:
|
||||
env (`gymnasium.Env`): A robot environment, built via `make_robot_env`.
|
||||
policy (`PreTrainedPolicy`): A policy exposing `select_action(obs) -> action`.
|
||||
n_episodes (`int`): Number of episodes to run.
|
||||
"""
|
||||
sum_reward_episode = []
|
||||
for _ in range(n_episodes):
|
||||
obs, _ = env.reset()
|
||||
@@ -54,6 +61,12 @@ def eval_policy(env, policy, n_episodes):
|
||||
|
||||
@parser.wrap()
|
||||
def main(cfg: TrainRLServerPipelineConfig):
|
||||
"""CLI entry point: load a pretrained policy and evaluate it for 10 episodes.
|
||||
|
||||
Args:
|
||||
cfg (`TrainRLServerPipelineConfig`): Parsed from the CLI. `cfg.env.pretrained_policy_name_or_path`
|
||||
selects the checkpoint to load; `cfg.dataset.repo_id` provides normalization stats.
|
||||
"""
|
||||
env_cfg = cfg.env
|
||||
env = make_robot_env(env_cfg)
|
||||
dataset_cfg = cfg.dataset
|
||||
|
||||
@@ -305,7 +305,9 @@ def make_robot_env(cfg: HILSerlRobotEnvConfig) -> tuple[gym.Env, Any]:
|
||||
"""Create robot environment from configuration.
|
||||
|
||||
Args:
|
||||
cfg: Environment configuration.
|
||||
cfg (`HILSerlRobotEnvConfig`): Environment configuration. `cfg.name == "gym_hil"` selects the
|
||||
GymHIL simulation environment; otherwise a real-robot `RobotEnv` is built from
|
||||
`cfg.robot`/`cfg.teleop`.
|
||||
|
||||
Returns:
|
||||
Tuple of (gym environment, teleoperator device).
|
||||
@@ -363,10 +365,13 @@ def make_processors(
|
||||
"""Create environment and action processors.
|
||||
|
||||
Args:
|
||||
env: Robot environment instance.
|
||||
teleop_device: Teleoperator device for intervention.
|
||||
cfg: Processor configuration.
|
||||
device: Target device for computations.
|
||||
env (`Env`): The environment returned by `make_robot_env`.
|
||||
teleop_device (`lerobot.teleoperators.teleoperator.Teleoperator | None`): The teleoperator
|
||||
device returned by `make_robot_env`, used to configure intervention-related processor
|
||||
steps. `None` for simulation environments.
|
||||
cfg (`HILSerlRobotEnvConfig`): Environment configuration; provides the reward classifier,
|
||||
gripper, and reset-behavior settings for the built processor steps.
|
||||
device (`str`, *optional*, defaults to `"cpu"`): Torch device the processors run on.
|
||||
|
||||
Returns:
|
||||
Tuple of (environment processor, action processor).
|
||||
@@ -536,20 +541,21 @@ def step_env_and_process_transition(
|
||||
env_processor: DataProcessorPipeline[EnvTransition, EnvTransition],
|
||||
action_processor: DataProcessorPipeline[EnvTransition, EnvTransition],
|
||||
) -> EnvTransition:
|
||||
"""
|
||||
Execute one step with processor pipeline.
|
||||
"""Execute one step with processor pipeline.
|
||||
|
||||
Args:
|
||||
env: The robot environment
|
||||
transition: Current transition state
|
||||
action: Action to execute
|
||||
env_processor: Environment processor
|
||||
action_processor: Action processor
|
||||
env (`Env`): The environment to step.
|
||||
transition (`EnvTransition`): The current transition; its observation is overwritten with the
|
||||
action processor's input before dispatch, then discarded.
|
||||
action (`Tensor`): The raw action to process and send to `env`.
|
||||
env_processor (`DataProcessorPipeline`): Post-processes the environment-produced transition
|
||||
(e.g. reward shaping, termination overrides).
|
||||
action_processor (`DataProcessorPipeline`): Pre-processes `action` before it reaches `env`
|
||||
(e.g. intervention overrides, gripper handling).
|
||||
|
||||
Returns:
|
||||
Processed transition with updated state.
|
||||
"""
|
||||
|
||||
# Create action transition
|
||||
transition[TransitionKey.ACTION] = action
|
||||
transition[TransitionKey.OBSERVATION] = (
|
||||
@@ -618,14 +624,16 @@ def control_loop(
|
||||
cfg: GymManipulatorConfig,
|
||||
) -> None:
|
||||
"""Main control loop for robot environment interaction.
|
||||
if cfg.mode == "record": then a dataset will be created and recorded
|
||||
|
||||
When `cfg.mode == "record"`, a dataset is created and recorded.
|
||||
|
||||
Args:
|
||||
env: The robot environment
|
||||
env_processor: Environment processor
|
||||
action_processor: Action processor
|
||||
teleop_device: Teleoperator device
|
||||
cfg: gym_manipulator configuration
|
||||
env (`Env`): The environment to control, built via `make_robot_env`.
|
||||
env_processor (`DataProcessorPipeline`): Post-processes environment-produced transitions.
|
||||
action_processor (`DataProcessorPipeline`): Pre-processes teleoperator actions before they
|
||||
reach `env`.
|
||||
teleop_device (`Teleoperator`): Teleoperator device driving the robot.
|
||||
cfg (`GymManipulatorConfig`): Control-loop configuration (mode, fps, episode/dataset settings).
|
||||
"""
|
||||
dt = 1.0 / cfg.env.fps
|
||||
|
||||
|
||||
@@ -31,18 +31,17 @@ from lerobot.utils.constants import OBS_STATE
|
||||
@dataclass
|
||||
@ProcessorStepRegistry.register("joint_velocity_processor")
|
||||
class JointVelocityProcessorStep(ObservationProcessorStep):
|
||||
"""
|
||||
Calculates and appends joint velocity information to the observation state.
|
||||
"""Calculates and appends joint velocity information to the observation state.
|
||||
|
||||
This step computes the velocity of each joint by calculating the finite
|
||||
difference between the current and the last observed joint positions. The
|
||||
resulting velocity vector is then concatenated to the original state vector.
|
||||
|
||||
Attributes:
|
||||
dt: The time step (delta time) in seconds between observations, used for
|
||||
calculating velocity.
|
||||
last_joint_positions: Stores the joint positions from the previous step
|
||||
to enable velocity calculation.
|
||||
**Attributes**:
|
||||
- **dt** (`float`) -- The time step (delta time) in seconds between observations, used for calculating
|
||||
velocity.
|
||||
- **last_joint_positions** (`torch.Tensor | None`) -- Stores the joint positions from the previous
|
||||
step to enable velocity calculation.
|
||||
"""
|
||||
|
||||
dt: float = 0.1
|
||||
@@ -50,8 +49,7 @@ class JointVelocityProcessorStep(ObservationProcessorStep):
|
||||
last_joint_positions: torch.Tensor | None = None
|
||||
|
||||
def observation(self, observation: dict) -> dict:
|
||||
"""
|
||||
Computes joint velocities and adds them to the observation state.
|
||||
"""Computes joint velocities and adds them to the observation state.
|
||||
|
||||
Args:
|
||||
observation: The input observation dictionary, expected to contain
|
||||
@@ -89,8 +87,7 @@ class JointVelocityProcessorStep(ObservationProcessorStep):
|
||||
return new_observation
|
||||
|
||||
def get_config(self) -> dict[str, Any]:
|
||||
"""
|
||||
Returns the configuration of the step for serialization.
|
||||
"""Returns the configuration of the step for serialization.
|
||||
|
||||
Returns:
|
||||
A dictionary containing the time step `dt`.
|
||||
@@ -106,8 +103,7 @@ class JointVelocityProcessorStep(ObservationProcessorStep):
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
"""
|
||||
Updates the `observation.state` feature to reflect the added velocities.
|
||||
"""Updates the `observation.state` feature to reflect the added velocities.
|
||||
|
||||
This method doubles the size of the first dimension of the `observation.state`
|
||||
shape to account for the concatenation of position and velocity vectors.
|
||||
@@ -132,22 +128,20 @@ class JointVelocityProcessorStep(ObservationProcessorStep):
|
||||
@dataclass
|
||||
@ProcessorStepRegistry.register("current_processor")
|
||||
class MotorCurrentProcessorStep(ObservationProcessorStep):
|
||||
"""
|
||||
Reads motor currents from a robot and appends them to the observation state.
|
||||
"""Reads motor currents from a robot and appends them to the observation state.
|
||||
|
||||
This step queries the robot's hardware interface to get the present current
|
||||
for each motor and concatenates this information to the existing state vector.
|
||||
|
||||
Attributes:
|
||||
robot: An instance of a `lerobot` Robot class that provides access to
|
||||
the hardware bus.
|
||||
**Attributes**:
|
||||
- **robot** (`Robot | None`) -- An instance of a `lerobot` Robot class that provides access to the
|
||||
hardware bus.
|
||||
"""
|
||||
|
||||
robot: Robot | None = None
|
||||
|
||||
def observation(self, observation: dict) -> dict:
|
||||
"""
|
||||
Fetches motor currents and adds them to the observation state.
|
||||
"""Fetches motor currents and adds them to the observation state.
|
||||
|
||||
Args:
|
||||
observation: The input observation dictionary.
|
||||
@@ -184,8 +178,7 @@ class MotorCurrentProcessorStep(ObservationProcessorStep):
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
"""
|
||||
Updates the `observation.state` feature to reflect the added motor currents.
|
||||
"""Updates the `observation.state` feature to reflect the added motor currents.
|
||||
|
||||
This method increases the size of the first dimension of the `observation.state`
|
||||
shape by the number of motors in the robot.
|
||||
|
||||
+85
-70
@@ -14,8 +14,7 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""
|
||||
Learner server runner for distributed HILSerl robot policy training.
|
||||
"""Learner server runner for distributed HILSerl robot policy training.
|
||||
|
||||
This script implements the learner component of the distributed HILSerl architecture.
|
||||
It initializes the policy network, maintains replay buffers, and updates
|
||||
@@ -121,6 +120,11 @@ from .trainer import RLTrainer
|
||||
|
||||
@parser.wrap()
|
||||
def train_cli(cfg: TrainRLServerPipelineConfig):
|
||||
"""CLI entry point for the HILSerl learner server.
|
||||
|
||||
Args:
|
||||
cfg (`TrainRLServerPipelineConfig`): Parsed from the CLI, forwarded to `train`.
|
||||
"""
|
||||
# Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
|
||||
require_package("grpcio", extra="hilserl", import_name="grpc")
|
||||
if not use_threads(cfg):
|
||||
@@ -136,14 +140,13 @@ def train_cli(cfg: TrainRLServerPipelineConfig):
|
||||
|
||||
|
||||
def train(cfg: TrainRLServerPipelineConfig, job_name: str | None = None):
|
||||
"""
|
||||
Main training function that initializes and runs the training process.
|
||||
"""Main training function that initializes and runs the training process.
|
||||
|
||||
Args:
|
||||
cfg (TrainRLServerPipelineConfig): The training configuration
|
||||
job_name (str | None, optional): Job name for logging. Defaults to None.
|
||||
cfg (`TrainRLServerPipelineConfig`): The training configuration.
|
||||
job_name (`str | None`, *optional*): Job name for logging. Defaults to `cfg.job_name` when
|
||||
unset.
|
||||
"""
|
||||
|
||||
cfg.validate()
|
||||
|
||||
if job_name is None:
|
||||
@@ -198,13 +201,12 @@ def start_learner_threads(
|
||||
wandb_logger: WandBLogger | None,
|
||||
shutdown_event: Any, # Event
|
||||
) -> None:
|
||||
"""
|
||||
Start the learner threads for training.
|
||||
"""Start the learner threads for training.
|
||||
|
||||
Args:
|
||||
cfg (TrainRLServerPipelineConfig): Training configuration
|
||||
wandb_logger (WandBLogger | None): Logger for metrics
|
||||
shutdown_event: Event to signal shutdown
|
||||
cfg (`TrainRLServerPipelineConfig`): Training configuration.
|
||||
wandb_logger (`WandBLogger | None`): Logger for metrics.
|
||||
shutdown_event (`Event`): Event signaling the learner and its background workers to stop.
|
||||
"""
|
||||
# Create multiprocessing queues
|
||||
transition_queue = Queue()
|
||||
@@ -275,9 +277,7 @@ def add_actor_information_and_train(
|
||||
interaction_message_queue: Queue,
|
||||
parameters_queue: Queue,
|
||||
):
|
||||
"""
|
||||
Handles data transfer from the actor to the learner, manages training updates,
|
||||
and logs training progress in an online reinforcement learning setup.
|
||||
"""Handles data transfer from the actor to the learner, manages training updates, and logs progress.
|
||||
|
||||
This function continuously:
|
||||
- Transfers transitions from the actor to the replay buffer.
|
||||
@@ -482,17 +482,18 @@ def start_learner(
|
||||
shutdown_event: Any, # Event
|
||||
cfg: TrainRLServerPipelineConfig,
|
||||
):
|
||||
"""
|
||||
Start the learner server for training.
|
||||
It will receive transitions and interaction messages from the actor server,
|
||||
and send policy parameters to the actor server.
|
||||
"""Start the learner server for training.
|
||||
|
||||
Receives transitions and interaction messages from the actor server, and sends policy parameters
|
||||
to the actor server.
|
||||
|
||||
Args:
|
||||
parameters_queue: Queue for sending policy parameters to the actor
|
||||
transition_queue: Queue for receiving transitions from the actor
|
||||
interaction_message_queue: Queue for receiving interaction messages from the actor
|
||||
shutdown_event: Event to signal shutdown
|
||||
cfg: Training configuration
|
||||
parameters_queue (`Queue`): Queue of serialized policy weights, drained and streamed to the
|
||||
actor by `LearnerService.StreamParameters`.
|
||||
transition_queue (`Queue`): Queue filled by `LearnerService.SendTransitions`.
|
||||
interaction_message_queue (`Queue`): Queue filled by `LearnerService.SendInteractions`.
|
||||
shutdown_event (`Event`): Event signaling this process/thread to stop.
|
||||
cfg (`TrainRLServerPipelineConfig`): Training configuration.
|
||||
"""
|
||||
if not use_threads(cfg):
|
||||
# Create a process-specific log file
|
||||
@@ -560,8 +561,7 @@ def save_training_checkpoint(
|
||||
preprocessor=None,
|
||||
postprocessor=None,
|
||||
) -> None:
|
||||
"""
|
||||
Save training checkpoint and associated data.
|
||||
"""Save training checkpoint and associated data.
|
||||
|
||||
This function performs the following steps:
|
||||
1. Creates a checkpoint directory with the current optimization step
|
||||
@@ -572,18 +572,26 @@ def save_training_checkpoint(
|
||||
6. If an offline replay buffer exists, saves it as a separate dataset
|
||||
|
||||
Args:
|
||||
cfg: Training configuration
|
||||
optimization_step: Current optimization step
|
||||
online_steps: Total number of online steps
|
||||
interaction_message: Dictionary containing interaction information
|
||||
policy: Policy model to save
|
||||
optimizers: Dictionary of optimizers
|
||||
replay_buffer: Replay buffer to save as dataset
|
||||
offline_replay_buffer: Optional offline replay buffer to save
|
||||
dataset_repo_id: Repository ID for dataset
|
||||
fps: Frames per second for dataset
|
||||
preprocessor: Optional preprocessor pipeline to save
|
||||
postprocessor: Optional postprocessor pipeline to save
|
||||
cfg (`TrainRLServerPipelineConfig`): Training configuration, saved alongside the checkpoint.
|
||||
optimization_step (`int`): Current optimization step; used to name the checkpoint directory.
|
||||
online_steps (`int`): Total number of online steps; used to size the checkpoint directory's
|
||||
zero-padded step number.
|
||||
interaction_message (`dict | None`): Latest interaction message; its `"Interaction step"`
|
||||
entry is saved for resuming training.
|
||||
policy (`Module`): Policy model to save.
|
||||
optimizers (`dict`): Dictionary of optimizers whose states are saved.
|
||||
replay_buffer (`ReplayBuffer`): Replay buffer to save as a dataset.
|
||||
algorithm (`lerobot.rl.algorithms.base.RLAlgorithm | None`, *optional*): Algorithm whose state
|
||||
dict (critic ensembles, temperature, etc.) should also be saved.
|
||||
offline_replay_buffer (`lerobot.rl.buffer.ReplayBuffer | None`, *optional*): Optional offline
|
||||
replay buffer, saved as a separate dataset when provided.
|
||||
dataset_repo_id (`str | None`, *optional*): Repository id used when converting the replay
|
||||
buffer(s) to a dataset.
|
||||
fps (`int`, *optional*, defaults to 30): Frames per second recorded on the saved dataset(s).
|
||||
preprocessor (`PolicyProcessorPipeline | None`, *optional*): Optional preprocessor pipeline to
|
||||
save alongside the policy.
|
||||
postprocessor (`PolicyProcessorPipeline | None`, *optional*): Optional postprocessor pipeline
|
||||
to save alongside the policy.
|
||||
"""
|
||||
logging.info(f"Checkpoint policy after step {optimization_step}")
|
||||
_num_digits = max(6, len(str(online_steps)))
|
||||
@@ -650,8 +658,7 @@ def save_training_checkpoint(
|
||||
|
||||
|
||||
def handle_resume_logic(cfg: TrainRLServerPipelineConfig) -> TrainRLServerPipelineConfig:
|
||||
"""
|
||||
Handle the resume logic for training.
|
||||
"""Handle the resume logic for training.
|
||||
|
||||
If resume is True:
|
||||
- Verifies that a checkpoint exists
|
||||
@@ -712,19 +719,19 @@ def load_training_state(
|
||||
algorithm: RLAlgorithm | None = None,
|
||||
device: str | torch.device = "cpu",
|
||||
):
|
||||
"""
|
||||
Loads the training state (optimizers, RNG, step + interaction step, and
|
||||
algorithm-owned tensors) from the most recent checkpoint.
|
||||
"""Loads the training state from the most recent checkpoint.
|
||||
|
||||
Restores optimizers, RNG state, the optimization/interaction step, and algorithm-owned tensors.
|
||||
|
||||
Args:
|
||||
cfg (TrainRLServerPipelineConfig): Training configuration; `cfg.resume` gates the load and
|
||||
cfg (`TrainRLServerPipelineConfig`): Training configuration; `cfg.resume` gates the load and
|
||||
`cfg.output_dir` locates the last checkpoint.
|
||||
optimizers (Optimizer | dict[str, Optimizer]): Optimizers to load state into.
|
||||
algorithm (RLAlgorithm | None, optional): Algorithm whose state dict should be restored.
|
||||
optimizers (`Optimizer | dict[str, Optimizer]`): Optimizers to load state into.
|
||||
algorithm (`RLAlgorithm | None`, *optional*): Algorithm whose state dict should be restored.
|
||||
Required for full main-equivalent resume; the policy itself is restored separately via
|
||||
`make_policy`. Defaults to None.
|
||||
device (str | torch.device, optional): Device on which to place loaded algorithm tensors.
|
||||
Defaults to "cpu".
|
||||
`make_policy`.
|
||||
device (`str | torch.device`, *optional*, defaults to `"cpu"`): Device on which to place
|
||||
loaded algorithm tensors.
|
||||
|
||||
Returns:
|
||||
tuple[int | None, int | None]: `(optimization_step, interaction_step)`, or `(None, None)`
|
||||
@@ -772,8 +779,7 @@ def load_training_state(
|
||||
|
||||
|
||||
def log_training_info(cfg: TrainRLServerPipelineConfig, policy: nn.Module) -> None:
|
||||
"""
|
||||
Log information about the training process.
|
||||
"""Log information about the training process.
|
||||
|
||||
Args:
|
||||
cfg (TrainRLServerPipelineConfig): Training configuration
|
||||
@@ -792,8 +798,7 @@ def log_training_info(cfg: TrainRLServerPipelineConfig, policy: nn.Module) -> No
|
||||
def initialize_replay_buffer(
|
||||
cfg: TrainRLServerPipelineConfig, device: str, storage_device: str
|
||||
) -> ReplayBuffer:
|
||||
"""
|
||||
Initialize a replay buffer, either empty or from a dataset if resuming.
|
||||
"""Initialize a replay buffer, either empty or from a dataset if resuming.
|
||||
|
||||
Args:
|
||||
cfg (TrainRLServerPipelineConfig): Training configuration
|
||||
@@ -837,8 +842,7 @@ def initialize_offline_replay_buffer(
|
||||
device: str,
|
||||
storage_device: str,
|
||||
) -> ReplayBuffer:
|
||||
"""
|
||||
Initialize an offline replay buffer from a dataset.
|
||||
"""Initialize an offline replay buffer from a dataset.
|
||||
|
||||
Args:
|
||||
cfg (TrainRLServerPipelineConfig): Training configuration
|
||||
@@ -875,6 +879,7 @@ def initialize_offline_replay_buffer(
|
||||
|
||||
|
||||
def use_threads(cfg: TrainRLServerPipelineConfig) -> bool:
|
||||
"""Whether the learner's background workers should run as threads instead of processes."""
|
||||
return cfg.policy.concurrency.learner == "threads"
|
||||
|
||||
|
||||
@@ -884,14 +889,14 @@ def check_nan_in_transition(
|
||||
next_state: torch.Tensor,
|
||||
raise_error: bool = False,
|
||||
) -> bool:
|
||||
"""
|
||||
Check for NaN values in transition data.
|
||||
"""Check for NaN values in transition data.
|
||||
|
||||
Args:
|
||||
observations: Dictionary of observation tensors
|
||||
actions: Action tensor
|
||||
next_state: Dictionary of next state tensors
|
||||
raise_error: If True, raises ValueError when NaN is detected
|
||||
observations (`Tensor`): Dictionary of observation tensors.
|
||||
actions (`Tensor`): Action tensor.
|
||||
next_state (`Tensor`): Dictionary of next-observation tensors.
|
||||
raise_error (`bool`, *optional*, defaults to `False`): Whether to raise a `ValueError` instead
|
||||
of just logging when a NaN is found.
|
||||
|
||||
Returns:
|
||||
bool: True if NaN values were detected, False otherwise
|
||||
@@ -925,6 +930,12 @@ def check_nan_in_transition(
|
||||
|
||||
|
||||
def push_actor_policy_to_queue(parameters_queue: Queue, algorithm: RLAlgorithm) -> None:
|
||||
"""Serialize `algorithm`'s current weights and enqueue them for the actor-facing gRPC stream.
|
||||
|
||||
Args:
|
||||
parameters_queue (`Queue`): Queue drained by `LearnerService.StreamParameters`.
|
||||
algorithm (`RLAlgorithm`): Source of the weights, via `get_weights`.
|
||||
"""
|
||||
logging.debug("[LEARNER] Pushing actor policy to the queue")
|
||||
|
||||
# Create a dictionary to hold all the state dicts
|
||||
@@ -958,11 +969,13 @@ def process_transitions(
|
||||
"""Process all available transitions from the queue.
|
||||
|
||||
Args:
|
||||
transition_queue: Queue for receiving transitions from the actor
|
||||
replay_buffer: Replay buffer to add transitions to
|
||||
offline_replay_buffer: Offline replay buffer to add transitions to
|
||||
dataset_repo_id: Repository ID for dataset
|
||||
shutdown_event: Event to signal shutdown
|
||||
transition_queue (`Queue`): Queue filled by `LearnerService.SendTransitions`.
|
||||
replay_buffer (`ReplayBuffer`): Buffer every non-NaN transition is added to.
|
||||
offline_replay_buffer (`ReplayBuffer`): Buffer intervention transitions are additionally added
|
||||
to, when `dataset_repo_id` is set.
|
||||
dataset_repo_id (`str | None`): When set, transitions tagged as interventions are also added
|
||||
to `offline_replay_buffer`.
|
||||
shutdown_event (`Event`): Event that stops the loop when set.
|
||||
"""
|
||||
while not transition_queue.empty() and not shutdown_event.is_set():
|
||||
transition_list = transition_queue.get()
|
||||
@@ -996,10 +1009,12 @@ def process_interaction_messages(
|
||||
"""Process all available interaction messages from the queue.
|
||||
|
||||
Args:
|
||||
interaction_message_queue: Queue for receiving interaction messages
|
||||
interaction_step_shift: Amount to shift interaction step by
|
||||
wandb_logger: Logger for tracking progress
|
||||
shutdown_event: Event to signal shutdown
|
||||
interaction_message_queue (`Queue`): Queue filled by `LearnerService.SendInteractions`.
|
||||
interaction_step_shift (`int`): Offset added to each message's `"Interaction step"` so it
|
||||
stays consistent with checkpointed state after a resume.
|
||||
wandb_logger (`lerobot.common.wandb_utils.WandBLogger | None`): Logger the message is
|
||||
forwarded to, when set.
|
||||
shutdown_event (`Event`): Event that stops the loop when set.
|
||||
|
||||
Returns:
|
||||
dict | None: The last interaction message processed, or None if none were processed
|
||||
|
||||
@@ -44,10 +44,10 @@ SHUTDOWN_TIMEOUT = 10
|
||||
|
||||
|
||||
class LearnerService(_ServicerBase):
|
||||
"""
|
||||
Implementation of the LearnerService gRPC service
|
||||
This service is used to send parameters to the Actor and receive transitions and interactions from the Actor
|
||||
check transport.proto for the gRPC service definition
|
||||
"""Implementation of the LearnerService gRPC service.
|
||||
|
||||
Sends policy parameters to the actor and receives transitions and interactions from it; see
|
||||
`transport.proto` for the gRPC service definition.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -59,6 +59,19 @@ class LearnerService(_ServicerBase):
|
||||
interaction_message_queue: Queue,
|
||||
queue_get_timeout: float = 0.001,
|
||||
):
|
||||
"""Create the servicer.
|
||||
|
||||
Args:
|
||||
shutdown_event (`Event`): Set to stop `StreamParameters`'s push loop.
|
||||
parameters_queue (`Queue`): Queue of serialized policy weights, drained and streamed to
|
||||
the actor by `StreamParameters`.
|
||||
seconds_between_pushes (`float`): Minimum interval between successive parameter pushes.
|
||||
transition_queue (`Queue`): Queue filled by `SendTransitions` with received transitions.
|
||||
interaction_message_queue (`Queue`): Queue filled by `SendInteractions` with received
|
||||
interaction messages.
|
||||
queue_get_timeout (`float`, *optional*, defaults to 0.001): Timeout used when polling
|
||||
`parameters_queue`.
|
||||
"""
|
||||
self.shutdown_event = shutdown_event
|
||||
self.parameters_queue = parameters_queue
|
||||
self.seconds_between_pushes = seconds_between_pushes
|
||||
@@ -69,6 +82,17 @@ class LearnerService(_ServicerBase):
|
||||
def StreamParameters( # noqa: N802
|
||||
self, request: "services_pb2.Empty", context: "grpc.ServicerContext"
|
||||
):
|
||||
"""GRPC server-streaming RPC: push the latest policy parameters to the actor.
|
||||
|
||||
Runs until `shutdown_event` is set, pushing at most once every `seconds_between_pushes`.
|
||||
|
||||
Args:
|
||||
request (`services_pb2.Empty`): Unused; required by the gRPC service signature.
|
||||
context (`grpc.ServicerContext`): gRPC call context.
|
||||
|
||||
Yields:
|
||||
Chunks of a `services_pb2.Parameters` message, produced by `send_bytes_in_chunks`.
|
||||
"""
|
||||
# TODO: authorize the request
|
||||
logging.info("[LEARNER] Received request to stream parameters from the Actor")
|
||||
|
||||
@@ -104,6 +128,16 @@ class LearnerService(_ServicerBase):
|
||||
return services_pb2.Empty()
|
||||
|
||||
def SendTransitions(self, request_iterator, _context: "grpc.ServicerContext"): # noqa: N802
|
||||
"""GRPC client-streaming RPC: receive transition chunks from the actor into `transition_queue`.
|
||||
|
||||
Args:
|
||||
request_iterator: Stream of `services_pb2.Transition` chunks sent by the actor's
|
||||
`transitions_stream`.
|
||||
_context (`grpc.ServicerContext`): gRPC call context.
|
||||
|
||||
Returns:
|
||||
services_pb2.Empty: Acknowledgement sent once the actor closes the stream.
|
||||
"""
|
||||
# TODO: authorize the request
|
||||
logging.info("[LEARNER] Received request to receive transitions from the Actor")
|
||||
|
||||
@@ -118,6 +152,16 @@ class LearnerService(_ServicerBase):
|
||||
return services_pb2.Empty()
|
||||
|
||||
def SendInteractions(self, request_iterator, _context: "grpc.ServicerContext"): # noqa: N802
|
||||
"""GRPC client-streaming RPC: receive interaction-message chunks into `interaction_message_queue`.
|
||||
|
||||
Args:
|
||||
request_iterator: Stream of `services_pb2.InteractionMessage` chunks sent by the actor's
|
||||
`interactions_stream`.
|
||||
_context (`grpc.ServicerContext`): gRPC call context.
|
||||
|
||||
Returns:
|
||||
services_pb2.Empty: Acknowledgement sent once the actor closes the stream.
|
||||
"""
|
||||
# TODO: authorize the request
|
||||
logging.info("[LEARNER] Received request to receive interactions from the Actor")
|
||||
|
||||
@@ -132,4 +176,5 @@ class LearnerService(_ServicerBase):
|
||||
return services_pb2.Empty()
|
||||
|
||||
def Ready(self, request: "services_pb2.Empty", context: "grpc.ServicerContext"): # noqa: N802
|
||||
"""GRPC health check: returns immediately, confirming the learner server is up."""
|
||||
return services_pb2.Empty()
|
||||
|
||||
@@ -23,6 +23,19 @@ from torch.multiprocessing import Queue
|
||||
|
||||
|
||||
def get_last_item_from_queue(queue: Queue, block=True, timeout: float = 0.1) -> Any:
|
||||
"""Drain `queue` and return only the most recently enqueued item.
|
||||
|
||||
Args:
|
||||
queue (`Queue`): A `torch.multiprocessing.Queue` to drain.
|
||||
block (`bool`, *optional*, defaults to `True`): Whether to block for up to `timeout` seconds
|
||||
waiting for a first item before draining. When `False`, returns `None` if the queue is
|
||||
currently empty.
|
||||
timeout (`float`, *optional*, defaults to 0.1): Seconds to wait for a first item when `block`
|
||||
is `True`.
|
||||
|
||||
Returns:
|
||||
Any: The most recent item, or `None` if the queue was (and stayed) empty.
|
||||
"""
|
||||
if block:
|
||||
try:
|
||||
item = queue.get(timeout=timeout)
|
||||
|
||||
@@ -28,6 +28,100 @@ from .algorithms.sac import SACAlgorithmConfig # noqa: F401
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class TrainRLServerPipelineConfig(TrainPipelineConfig):
|
||||
"""Top-level config for the actor/learner distributed RL training server.
|
||||
|
||||
Extends [`~configs.train.TrainPipelineConfig`] with an optional (rather than required) offline
|
||||
`dataset` and the RL-specific algorithm/data-mixing fields below.
|
||||
|
||||
Args:
|
||||
env (`lerobot.envs.configs.EnvConfig | None`, *optional*):
|
||||
Simulation environment configuration, used for `env_eval_freq` evaluation rollouts.
|
||||
policy (`lerobot.configs.policies.PreTrainedConfig | None`, *optional*):
|
||||
The actor policy configuration.
|
||||
reward_model (`lerobot.configs.rewards.RewardModelConfig | None`, *optional*):
|
||||
Reward model configuration, when training a reward model instead of a policy.
|
||||
output_dir (`pathlib.Path | None`, *optional*):
|
||||
Directory to save run outputs to. Reusing the same value across runs overwrites its
|
||||
contents unless `resume` is `True`.
|
||||
job_name (`str | None`, *optional*):
|
||||
Name used for logging and checkpoint directory naming.
|
||||
resume (`bool`, *optional*, defaults to `False`):
|
||||
Whether to resume a previous run from `--config_path`'s checkpoint.
|
||||
seed (`int | None`, *optional*, defaults to 1000):
|
||||
Random seed for model initialization, dataset shuffling, and evaluation environments.
|
||||
cudnn_deterministic (`bool`, *optional*, defaults to `False`):
|
||||
Whether to use deterministic cuDNN algorithms for reproducibility. Disables
|
||||
`cudnn.benchmark`, which may reduce training speed.
|
||||
num_workers (`int`, *optional*, defaults to 4):
|
||||
Number of dataloader worker processes.
|
||||
batch_size (`int`, *optional*, defaults to 8):
|
||||
Offline-dataset dataloader batch size.
|
||||
prefetch_factor (`int`, *optional*, defaults to 4):
|
||||
Number of batches prefetched per dataloader worker.
|
||||
persistent_workers (`bool`, *optional*, defaults to `True`):
|
||||
Whether dataloader workers stay alive between epochs.
|
||||
dataloader_multiprocessing_context (`str | None`, *optional*, defaults to `"spawn"`):
|
||||
DataLoader worker start method. `None` uses Python's platform default.
|
||||
steps (`int`, *optional*, defaults to 100000):
|
||||
Total number of training steps.
|
||||
env_eval_freq (`int`, *optional*, defaults to 20000):
|
||||
Run the policy in the simulation environment every N steps to measure reward/success.
|
||||
`0` disables environment evaluation.
|
||||
log_freq (`int`, *optional*, defaults to 200):
|
||||
Log training metrics every N steps.
|
||||
eval_steps (`int`, *optional*, defaults to 0):
|
||||
Compute eval loss on held-out episodes every N steps. `0` disables it.
|
||||
max_eval_samples (`int`, *optional*, defaults to 0):
|
||||
Cap on total eval samples, split uniformly across tasks. `0` uses all held-out data.
|
||||
tolerance_s (`float`, *optional*, defaults to 0.0001):
|
||||
Maximum timestamp tolerance, in seconds, when loading dataset frames.
|
||||
save_checkpoint (`bool`, *optional*, defaults to `True`):
|
||||
Whether to save training checkpoints at all.
|
||||
save_freq (`int`, *optional*, defaults to 20000):
|
||||
Save a checkpoint every N training steps, and after the last step. A non-positive value
|
||||
disables periodic saving, keeping only the final checkpoint.
|
||||
checkpoint_format (`CheckpointFormat`, *optional*, defaults to `CheckpointFormat.SAFETENSORS`):
|
||||
Model-artifact format inside checkpoints.
|
||||
use_policy_training_preset (`bool`, *optional*, defaults to `True`):
|
||||
Whether to use the policy's own recommended optimizer/scheduler preset when `optimizer`/
|
||||
`scheduler` are unset.
|
||||
optimizer (`lerobot.optim.optimizers.OptimizerConfig | None`, *optional*):
|
||||
Optimizer configuration override.
|
||||
scheduler (`lerobot.optim.schedulers.LRSchedulerConfig | None`, *optional*):
|
||||
Learning-rate scheduler configuration override.
|
||||
parallelism (`ParallelismConfig`, *optional*):
|
||||
Process topology: `dp_replicate`/`dp_shard` (HSDP) and context-parallel degree.
|
||||
accelerator (`AcceleratorConfig`, *optional*):
|
||||
Execution runtime handed to the Accelerator: mixed precision, gradient accumulation,
|
||||
FSDP/DDP tuning knobs, compile & activation-checkpointing.
|
||||
eval (`EvalConfig`, *optional*):
|
||||
Simulation-environment evaluation configuration (number of episodes, batch size).
|
||||
wandb (`WandBConfig`, *optional*):
|
||||
Weights & Biases logging configuration.
|
||||
peft (`lerobot.configs.default.PeftConfig | None`, *optional*):
|
||||
PEFT (e.g. LoRA) adapter configuration for parameter-efficient fine-tuning.
|
||||
job (`JobConfig`, *optional*):
|
||||
Where to run training: local (default) or an HF Jobs flavor.
|
||||
save_checkpoint_to_hub (`bool`, *optional*, defaults to `False`):
|
||||
Whether to push each saved checkpoint to the Hub as it is written, not just the final
|
||||
model.
|
||||
sample_weighting (`lerobot.utils.sample_weighting.SampleWeightingConfig | None`, *optional*):
|
||||
Sample weighting configuration (e.g. for RA-BC training).
|
||||
rename_map (`dict`, *optional*):
|
||||
Mapping to override observation image/state key names.
|
||||
dataset (`DatasetConfig | None`, *optional*):
|
||||
Optional offline dataset config. Unlike imitation-learning training, RL doesn't require an
|
||||
offline dataset — data comes from the online replay buffer.
|
||||
algorithm (`RLAlgorithmConfig | None`, *optional*):
|
||||
RL algorithm configuration. Defaults to a SAC config (with `policy_config` populated from
|
||||
`self.policy`) in `validate` when unset.
|
||||
mixer (`str`, *optional*, defaults to `"online_offline"`):
|
||||
Data mixer strategy name. Currently only `"online_offline"` is supported.
|
||||
online_ratio (`float`, *optional*, defaults to 0.5):
|
||||
Fraction of each training batch sampled from the online replay buffer when using
|
||||
`OnlineOfflineMixer`; the remainder comes from the offline dataset.
|
||||
"""
|
||||
|
||||
# NOTE: In RL, we don't need an offline dataset
|
||||
# TODO: Make `TrainPipelineConfig.dataset` optional
|
||||
dataset: DatasetConfig | None = None # type: ignore[assignment] # because the parent class has made it's type non-optional
|
||||
@@ -41,6 +135,11 @@ class TrainRLServerPipelineConfig(TrainPipelineConfig):
|
||||
online_ratio: float = 0.5
|
||||
|
||||
def validate(self) -> None:
|
||||
"""See [`~configs.train.TrainPipelineConfig.validate`].
|
||||
|
||||
Additionally defaults `algorithm` to a SAC config and populates its `policy_config` from
|
||||
`self.policy` when unset.
|
||||
"""
|
||||
super().validate()
|
||||
|
||||
if self.algorithm is None:
|
||||
|
||||
@@ -38,6 +38,16 @@ class RLTrainer:
|
||||
*,
|
||||
preprocessor: Any | None = None,
|
||||
):
|
||||
"""Build the trainer and its optimizers.
|
||||
|
||||
Args:
|
||||
algorithm (`RLAlgorithm`): The RL algorithm to train. `make_optimizers_and_scheduler` is
|
||||
called on it immediately.
|
||||
data_mixer (`DataMixer`): Data source the training-batch iterator is built from.
|
||||
batch_size (`int`): Batch size requested from `data_mixer` on each training step.
|
||||
preprocessor (`Any | None`, *optional*): When set, each sampled batch is passed through
|
||||
`preprocess_rl_batch` before reaching the algorithm.
|
||||
"""
|
||||
self.algorithm = algorithm
|
||||
self.data_mixer = data_mixer
|
||||
self.batch_size = batch_size
|
||||
@@ -90,12 +100,15 @@ class _PreprocessedIterator:
|
||||
__slots__ = ("_raw", "_preprocessor")
|
||||
|
||||
def __init__(self, raw_iterator: Iterator[BatchType], preprocessor: Any) -> None:
|
||||
"""Wrap `raw_iterator`, applying `preprocessor` to each yielded batch."""
|
||||
self._raw = raw_iterator
|
||||
self._preprocessor = preprocessor
|
||||
|
||||
def __iter__(self) -> _PreprocessedIterator:
|
||||
"""Return `self` (this object is its own iterator)."""
|
||||
return self
|
||||
|
||||
def __next__(self) -> BatchType:
|
||||
"""Return the next preprocessed batch from the wrapped iterator."""
|
||||
batch = next(self._raw)
|
||||
return preprocess_rl_batch(self._preprocessor, batch)
|
||||
|
||||
@@ -29,14 +29,18 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BiOpenArmFollower(BimanualMixin, Robot):
|
||||
"""
|
||||
Bimanual OpenArm Follower Arms
|
||||
"""
|
||||
"""A bimanual pair of OpenArm follower arms driven as one robot."""
|
||||
|
||||
config_class = BiOpenArmFollowerConfig
|
||||
name = "bi_openarm_follower"
|
||||
|
||||
def __init__(self, config: BiOpenArmFollowerConfig):
|
||||
"""Build the robot from its configuration.
|
||||
|
||||
Args:
|
||||
config (`BiOpenArmFollowerConfig`):
|
||||
The robot's configuration. Its `port` and `cameras` determine what is connected.
|
||||
"""
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
|
||||
@@ -114,19 +118,43 @@ class BiOpenArmFollower(BimanualMixin, Robot):
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
"""The values this robot reports, and their types or shapes.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys as returned by [`~robots.Robot.get_observation`], mapped to a scalar type for
|
||||
proprioceptive values or to a `(height, width, channels)` shape for images.
|
||||
"""
|
||||
return {**self._motors_ft, **self._cameras_ft}
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
"""The values this robot accepts, and their types.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys accepted by [`~robots.Robot.send_action`], mapped to their type.
|
||||
"""
|
||||
return self._motors_ft
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
"""Assign each motor its bus ID, one at a time.
|
||||
|
||||
Run this once when building the robot. Interactive: prompts you to connect the controller board to a
|
||||
single motor at a time.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Motor ID configuration is typically done via manufacturer tools for CAN motors."
|
||||
)
|
||||
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
"""Read the robot's current state and a frame from each camera.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: Keys matching [`~robots.Robot.observation_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
obs_dict: RobotObservation = {}
|
||||
|
||||
# Add "left_" prefix to per-arm keys; keep top-level camera keys unprefixed.
|
||||
@@ -146,6 +174,23 @@ class BiOpenArmFollower(BimanualMixin, Robot):
|
||||
custom_kp: dict[str, float] | None = None,
|
||||
custom_kd: dict[str, float] | None = None,
|
||||
) -> RobotAction:
|
||||
"""Command both arms to move towards a target configuration.
|
||||
|
||||
Args:
|
||||
action (`dict[str, Any]`):
|
||||
Target values, keyed as in [`~robots.Robot.action_features`], i.e. prefixed `left_` and
|
||||
`right_`.
|
||||
custom_kp (`dict[str, float]`, *optional*):
|
||||
Per-motor proportional gains for this step only. Defaults to each arm's `position_kp`.
|
||||
custom_kd (`dict[str, float]`, *optional*):
|
||||
Per-motor derivative gains for this step only. Defaults to each arm's `position_kd`.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The action actually sent, which may be clipped by `max_relative_target`.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
# Remove "left_" prefix
|
||||
left_action = {
|
||||
key.removeprefix("left_"): value for key, value in action.items() if key.startswith("left_")
|
||||
|
||||
@@ -25,7 +25,28 @@ from ..openarm_follower import OpenArmFollowerConfigBase
|
||||
@RobotConfig.register_subclass("bi_openarm_follower")
|
||||
@dataclass(kw_only=True)
|
||||
class BiOpenArmFollowerConfig(RobotConfig):
|
||||
"""Configuration class for Bi OpenArm Follower robots."""
|
||||
"""Configuration for a bimanual pair of OpenArm follower arms.
|
||||
|
||||
The two arms are configured independently, then driven as one robot: observation and action keys from
|
||||
each arm are prefixed with `left_` and `right_`.
|
||||
|
||||
Calibration is per arm, taken from each arm config's own settings.
|
||||
|
||||
Args:
|
||||
id (`str`, *optional*, defaults to `"bi_openarm_follower"`):
|
||||
Identifier for the pair as a whole.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Unused at this level; each arm calibrates through its own config.
|
||||
left_arm_config (`OpenArmFollowerConfigBase`):
|
||||
Configuration for the left arm, including its own CAN interface. Set its `side` to `"left"` so
|
||||
the correct joint limits apply.
|
||||
right_arm_config (`OpenArmFollowerConfigBase`):
|
||||
Configuration for the right arm, including its own CAN interface. Set its `side` to `"right"`.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
Cameras not attached to either arm, such as an overhead view. These keys appear in
|
||||
observations unchanged, whereas cameras declared on an arm config are prefixed with that
|
||||
arm's side.
|
||||
"""
|
||||
|
||||
id: str | None = "bi_openarm_follower"
|
||||
|
||||
|
||||
@@ -39,6 +39,12 @@ class BiRebotB601Follower(BimanualMixin, Robot):
|
||||
name = "bi_rebot_b601_follower"
|
||||
|
||||
def __init__(self, config: BiRebotB601FollowerConfig):
|
||||
"""Build the robot from its configuration.
|
||||
|
||||
Args:
|
||||
config (`BiRebotB601FollowerConfig`):
|
||||
The robot's configuration. Its `port` and `cameras` determine what is connected.
|
||||
"""
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
|
||||
@@ -120,14 +126,33 @@ class BiRebotB601Follower(BimanualMixin, Robot):
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
"""The values this robot reports, and their types or shapes.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys as returned by [`~robots.Robot.get_observation`], mapped to a scalar type for
|
||||
proprioceptive values or to a `(height, width, channels)` shape for images.
|
||||
"""
|
||||
return {**self._motors_ft, **self._cameras_ft}
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
"""The values this robot accepts, and their types.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys accepted by [`~robots.Robot.send_action`], mapped to their type.
|
||||
"""
|
||||
return self._motors_ft
|
||||
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
"""Read the robot's current state and a frame from each camera.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: Keys matching [`~robots.Robot.observation_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
obs_dict: RobotObservation = {}
|
||||
for k, v in self.left_arm.get_observation().items():
|
||||
obs_dict[k if k in self._top_level_cam_keys else f"left_{k}"] = v
|
||||
@@ -137,6 +162,18 @@ class BiRebotB601Follower(BimanualMixin, Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def send_action(self, action: RobotAction) -> RobotAction:
|
||||
"""Command the robot to move towards a target configuration.
|
||||
|
||||
Args:
|
||||
action (`dict[str, Any]`):
|
||||
Target values, keyed as in [`~robots.Robot.action_features`].
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The action actually sent, which may be clipped by `max_relative_target`.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
left_action = {
|
||||
key.removeprefix("left_"): value for key, value in action.items() if key.startswith("left_")
|
||||
}
|
||||
|
||||
@@ -25,7 +25,27 @@ from ..rebot_b601_follower import RebotB601FollowerConfig
|
||||
@RobotConfig.register_subclass("bi_rebot_b601_follower")
|
||||
@dataclass
|
||||
class BiRebotB601FollowerConfig(RobotConfig):
|
||||
"""Configuration class for the bimanual reBot B601-DM follower robot."""
|
||||
"""Configuration for a bimanual pair of reBot B601-DM follower arms.
|
||||
|
||||
The two arms are configured independently, then driven as one robot: observation and action keys from
|
||||
each arm are prefixed with `left_` and `right_`.
|
||||
|
||||
Calibration is per arm, taken from each arm config's own `id` and `calibration_dir`.
|
||||
|
||||
Args:
|
||||
left_arm_config (`RebotB601FollowerConfig`):
|
||||
Configuration for the left arm, including its own `port` and CAN settings.
|
||||
right_arm_config (`RebotB601FollowerConfig`):
|
||||
Configuration for the right arm, including its own `port` and CAN settings.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
Cameras not attached to either arm, such as an overhead view. These keys appear in
|
||||
observations unchanged, whereas cameras declared on an arm config are prefixed with that
|
||||
arm's side.
|
||||
id (`str`, *optional*):
|
||||
Identifier for the pair as a whole.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Unused at this level; each arm calibrates through its own config.
|
||||
"""
|
||||
|
||||
left_arm_config: RebotB601FollowerConfig
|
||||
right_arm_config: RebotB601FollowerConfig
|
||||
|
||||
@@ -29,14 +29,18 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BiSOFollower(BimanualMixin, Robot):
|
||||
"""
|
||||
[Bimanual SO Follower Arms](https://github.com/TheRobotStudio/SO-ARM100) designed by TheRobotStudio
|
||||
"""
|
||||
"""A bimanual pair of [SO follower arms](https://github.com/TheRobotStudio/SO-ARM100) by TheRobotStudio."""
|
||||
|
||||
config_class = BiSOFollowerConfig
|
||||
name = "bi_so_follower"
|
||||
|
||||
def __init__(self, config: BiSOFollowerConfig):
|
||||
"""Build the robot from its configuration.
|
||||
|
||||
Args:
|
||||
config (`BiSOFollowerConfig`):
|
||||
The robot's configuration. Its `port` and `cameras` determine what is connected.
|
||||
"""
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
|
||||
@@ -107,18 +111,42 @@ class BiSOFollower(BimanualMixin, Robot):
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
"""The values this robot reports, and their types or shapes.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys as returned by [`~robots.Robot.get_observation`], mapped to a scalar type for
|
||||
proprioceptive values or to a `(height, width, channels)` shape for images.
|
||||
"""
|
||||
return {**self._motors_ft, **self._cameras_ft}
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
"""The values this robot accepts, and their types.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys accepted by [`~robots.Robot.send_action`], mapped to their type.
|
||||
"""
|
||||
return self._motors_ft
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
"""Assign each motor its bus ID, one at a time.
|
||||
|
||||
Run this once when building the robot. Interactive: prompts you to connect the controller board to a
|
||||
single motor at a time.
|
||||
"""
|
||||
self.left_arm.setup_motors()
|
||||
self.right_arm.setup_motors()
|
||||
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
"""Read the robot's current state and a frame from each camera.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: Keys matching [`~robots.Robot.observation_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
obs_dict: RobotObservation = {}
|
||||
|
||||
# Add "left_" prefix to per-arm keys; keep top-level camera keys unprefixed.
|
||||
@@ -134,6 +162,18 @@ class BiSOFollower(BimanualMixin, Robot):
|
||||
@check_if_not_connected
|
||||
def send_action(self, action: RobotAction) -> RobotAction:
|
||||
# Remove "left_" prefix
|
||||
"""Command the robot to move towards a target configuration.
|
||||
|
||||
Args:
|
||||
action (`dict[str, Any]`):
|
||||
Target values, keyed as in [`~robots.Robot.action_features`].
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The action actually sent, which may be clipped by `max_relative_target`.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
left_action = {
|
||||
key.removeprefix("left_"): value for key, value in action.items() if key.startswith("left_")
|
||||
}
|
||||
|
||||
@@ -25,7 +25,27 @@ from ..so_follower import SOFollowerConfig
|
||||
@RobotConfig.register_subclass("bi_so_follower")
|
||||
@dataclass
|
||||
class BiSOFollowerConfig(RobotConfig):
|
||||
"""Configuration class for Bi SO Follower robots."""
|
||||
"""Configuration for a bimanual pair of SO follower arms.
|
||||
|
||||
The two arms are configured independently, then driven as one robot: observation and action keys from
|
||||
each arm are prefixed with `left_` and `right_`.
|
||||
|
||||
Calibration is per arm, taken from each arm config's own `id` and `calibration_dir`.
|
||||
|
||||
Args:
|
||||
left_arm_config (`SOFollowerConfig`):
|
||||
Configuration for the left arm, including its own `port`.
|
||||
right_arm_config (`SOFollowerConfig`):
|
||||
Configuration for the right arm, including its own `port`.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
Cameras not attached to either arm, such as an overhead view. These keys appear in
|
||||
observations unchanged, whereas cameras declared on an arm config are prefixed with that
|
||||
arm's side.
|
||||
id (`str`, *optional*):
|
||||
Identifier for the pair as a whole.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Unused at this level; each arm calibrates through its own config.
|
||||
"""
|
||||
|
||||
left_arm_config: SOFollowerConfig
|
||||
right_arm_config: SOFollowerConfig
|
||||
|
||||
@@ -21,12 +21,33 @@ import draccus
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class RobotConfig(draccus.ChoiceRegistry, abc.ABC):
|
||||
"""Base configuration shared by every robot.
|
||||
|
||||
Concrete robots subclass this and register themselves with
|
||||
`@RobotConfig.register_subclass("name")`, which is what makes `--robot.type=name` work on the command
|
||||
line. Subclasses inherit the two fields below and must document them alongside their own.
|
||||
|
||||
Args:
|
||||
id (`str`, *optional*):
|
||||
Identifier for this particular unit, used to tell apart several robots of the same type. It
|
||||
also names the calibration file, so keep it stable for a given piece of hardware.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Where to read and write the calibration file. Defaults to a per-robot directory under the
|
||||
LeRobot calibration home.
|
||||
"""
|
||||
|
||||
# Allows to distinguish between different robots of the same type
|
||||
id: str | None = None
|
||||
# Directory to store calibration file
|
||||
calibration_dir: Path | None = None
|
||||
|
||||
def __post_init__(self):
|
||||
"""Validate that every configured camera specifies the fields a robot requires.
|
||||
|
||||
Raises:
|
||||
ValueError: If a camera does not set `width`, `height` and `fps`. A robot records frames at a
|
||||
fixed shape, so these cannot be left to the driver's defaults.
|
||||
"""
|
||||
if hasattr(self, "cameras") and self.cameras:
|
||||
for _, config in self.cameras.items():
|
||||
for attr in ["width", "height", "fps"]:
|
||||
@@ -37,4 +58,10 @@ class RobotConfig(draccus.ChoiceRegistry, abc.ABC):
|
||||
|
||||
@property
|
||||
def type(self) -> str:
|
||||
"""The registered name of this robot type.
|
||||
|
||||
Returns:
|
||||
`str`: The name passed to `@RobotConfig.register_subclass`, e.g. `"so101_follower"`. This is
|
||||
what `make_robot_from_config` dispatches on and what a user writes as `--robot.type=...`.
|
||||
"""
|
||||
return self.get_choice_name(self.__class__)
|
||||
|
||||
@@ -23,13 +23,18 @@ from ..config import RobotConfig
|
||||
@RobotConfig.register_subclass("earthrover_mini_plus")
|
||||
@dataclass
|
||||
class EarthRoverMiniPlusConfig(RobotConfig):
|
||||
"""Configuration for EarthRover Mini Plus robot using Frodobots SDK.
|
||||
"""Configuration for the EarthRover Mini Plus rover.
|
||||
|
||||
This robot uses cloud-based control via the Frodobots SDK HTTP API.
|
||||
Camera frames are accessed directly through SDK HTTP endpoints.
|
||||
This robot is driven over the cloud through the Frodobots SDK's HTTP API rather than a local bus, so
|
||||
there is no serial port and no LeRobot calibration file. Camera frames come from SDK HTTP endpoints.
|
||||
|
||||
Attributes:
|
||||
sdk_url: URL of the Frodobots SDK server (default: http://localhost:8000)
|
||||
Args:
|
||||
sdk_url (`str`, *optional*, defaults to `"http://localhost:8000"`):
|
||||
Base URL of the Frodobots SDK server. Commands and camera frames both go through it.
|
||||
id (`str`, *optional*):
|
||||
Identifier for this particular rover.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Unused: the rover exposes no calibration.
|
||||
"""
|
||||
|
||||
sdk_url: str = "http://localhost:8000"
|
||||
|
||||
@@ -70,8 +70,7 @@ OBS_WHEEL_RPM_3 = "wheel_rpm_3"
|
||||
|
||||
|
||||
class EarthRoverMiniPlus(Robot):
|
||||
"""
|
||||
EarthRover Mini Plus robot controlled via Frodobots SDK HTTP API.
|
||||
"""EarthRover Mini Plus robot controlled via Frodobots SDK HTTP API.
|
||||
|
||||
This robot uses cloud-based control through the Frodobots SDK instead of direct
|
||||
hardware connection. Cameras stream via WebRTC through Agora cloud, and control
|
||||
@@ -82,9 +81,9 @@ class EarthRoverMiniPlus(Robot):
|
||||
- Linear and angular velocity control
|
||||
- Battery and orientation telemetry
|
||||
|
||||
Attributes:
|
||||
config: Robot configuration
|
||||
sdk_base_url: URL of the Frodobots SDK server (default: http://localhost:8000)
|
||||
**Attributes**:
|
||||
- **config** -- Robot configuration
|
||||
- **sdk_base_url** -- URL of the Frodobots SDK server (default: http://localhost:8000)
|
||||
"""
|
||||
|
||||
config_class = EarthRoverMiniPlusConfig
|
||||
@@ -130,7 +129,6 @@ class EarthRoverMiniPlus(Robot):
|
||||
DeviceAlreadyConnectedError: If robot is already connected
|
||||
DeviceNotConnectedError: If cannot connect to SDK server
|
||||
"""
|
||||
|
||||
# Verify SDK is running and accessible
|
||||
try:
|
||||
response = requests.get(f"{self.sdk_base_url}/data", timeout=10.0)
|
||||
@@ -280,7 +278,6 @@ class EarthRoverMiniPlus(Robot):
|
||||
Robot telemetry is retrieved from /data endpoint.
|
||||
All SDK values are normalized to appropriate ranges for dataset recording.
|
||||
"""
|
||||
|
||||
observation = {}
|
||||
|
||||
# Get camera images from SDK
|
||||
@@ -370,7 +367,6 @@ class EarthRoverMiniPlus(Robot):
|
||||
Raises:
|
||||
DeviceNotConnectedError: If robot is not connected
|
||||
"""
|
||||
|
||||
# Stop the robot before disconnecting
|
||||
try:
|
||||
self._send_command_to_sdk(0.0, 0.0)
|
||||
|
||||
@@ -24,6 +24,27 @@ from ..config import RobotConfig
|
||||
@RobotConfig.register_subclass("hope_jr_hand")
|
||||
@dataclass
|
||||
class HopeJrHandConfig(RobotConfig):
|
||||
"""Configuration for one Hope Jr hand.
|
||||
|
||||
Each hand is a separate robot, so a two-handed setup uses two of these with different `side` and
|
||||
`port` values.
|
||||
|
||||
Args:
|
||||
port (`str`):
|
||||
Serial port the hand is connected to. Run `lerobot-find-port` to identify it.
|
||||
side (`str`):
|
||||
Which hand this is, `"left"` or `"right"`. Determines the motor layout, so it must match the
|
||||
hardware.
|
||||
disable_torque_on_disconnect (`bool`, *optional*, defaults to `True`):
|
||||
Whether to release the motors on disconnect.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
Cameras to read alongside the hand's joint positions.
|
||||
id (`str`, *optional*):
|
||||
Identifier for this particular hand; also names its calibration file.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
|
||||
"""
|
||||
|
||||
port: str # Port to connect to the hand
|
||||
side: str # "left" / "right"
|
||||
|
||||
@@ -32,6 +53,12 @@ class HopeJrHandConfig(RobotConfig):
|
||||
cameras: dict[str, CameraConfig] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self):
|
||||
"""Validate the camera settings and the hand side.
|
||||
|
||||
Raises:
|
||||
ValueError: If `side` is not `"left"` or `"right"`, or if a camera omits `width`, `height` or
|
||||
`fps`.
|
||||
"""
|
||||
super().__post_init__()
|
||||
if self.side not in ["right", "left"]:
|
||||
raise ValueError(self.side)
|
||||
@@ -40,6 +67,26 @@ class HopeJrHandConfig(RobotConfig):
|
||||
@RobotConfig.register_subclass("hope_jr_arm")
|
||||
@dataclass
|
||||
class HopeJrArmConfig(RobotConfig):
|
||||
"""Configuration for one Hope Jr arm.
|
||||
|
||||
Args:
|
||||
port (`str`):
|
||||
Serial port the arm is connected to. Run `lerobot-find-port` to identify it.
|
||||
disable_torque_on_disconnect (`bool`, *optional*, defaults to `True`):
|
||||
Whether to release the motors on disconnect. Leave `True` unless the arm is holding a load it
|
||||
must not drop.
|
||||
max_relative_target (`float | dict[str, float]`, *optional*):
|
||||
Caps how far a single action may move the arm from its present position, as a safety limit. A
|
||||
scalar applies to every motor; a dict maps motor name to a per-motor cap. `None` disables
|
||||
clipping.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
Cameras to read alongside the arm's joint positions.
|
||||
id (`str`, *optional*):
|
||||
Identifier for this particular arm; also names its calibration file.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
|
||||
"""
|
||||
|
||||
port: str # Port to connect to the hand
|
||||
disable_torque_on_disconnect: bool = True
|
||||
|
||||
|
||||
@@ -35,10 +35,22 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class HopeJrArm(Robot):
|
||||
"""One arm of the Hope Jr humanoid.
|
||||
|
||||
The arm and the hand are separate robots; pair this with [`~robots.hope_jr.HopeJrHand`] for a full
|
||||
limb. See [`~robots.Robot`] for the contract every method here implements.
|
||||
"""
|
||||
|
||||
config_class = HopeJrArmConfig
|
||||
name = "hope_jr_arm"
|
||||
|
||||
def __init__(self, config: HopeJrArmConfig):
|
||||
"""Build the robot from its configuration.
|
||||
|
||||
Args:
|
||||
config (`HopeJrArmConfig`):
|
||||
The robot's configuration. Its `port` and `cameras` determine what is connected.
|
||||
"""
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
self.bus = FeetechMotorsBus(
|
||||
@@ -77,23 +89,47 @@ class HopeJrArm(Robot):
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
"""The values this robot reports, and their types or shapes.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys as returned by [`~robots.Robot.get_observation`], mapped to a scalar type for
|
||||
proprioceptive values or to a `(height, width, channels)` shape for images.
|
||||
"""
|
||||
return {**self._motors_ft, **self._cameras_ft}
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
"""The values this robot accepts, and their types.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys accepted by [`~robots.Robot.send_action`], mapped to their type.
|
||||
"""
|
||||
return self._motors_ft
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Whether every device this robot uses is connected.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` only when the robot and all its cameras are connected.
|
||||
"""
|
||||
return self.bus.is_connected and all(cam.is_connected for cam in self.cameras.values())
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""
|
||||
We assume that at connection time, arm is in a rest position,
|
||||
and torque can be safely disabled to run calibration.
|
||||
"""
|
||||
"""Connect the motor bus and cameras, calibrating and configuring the arm.
|
||||
|
||||
> [!WARNING]
|
||||
> The arm is assumed to be at rest when this is called, because torque is disabled to run
|
||||
> calibration.
|
||||
|
||||
Args:
|
||||
calibrate (`bool`, *optional*, defaults to `True`):
|
||||
Whether to run calibration if the arm is not already calibrated.
|
||||
|
||||
Raises:
|
||||
DeviceAlreadyConnectedError: If the robot is already connected.
|
||||
"""
|
||||
self.bus.connect(handshake=False)
|
||||
if not self.is_calibrated and calibrate:
|
||||
self.calibrate()
|
||||
@@ -107,9 +143,18 @@ class HopeJrArm(Robot):
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
"""Whether the robot is calibrated.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` when no calibration is needed before use.
|
||||
"""
|
||||
return self.bus.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
"""Calibrate the robot and store the result.
|
||||
|
||||
Interactive: prompts on stdin and asks you to move the robot through the required positions.
|
||||
"""
|
||||
groups = {
|
||||
"all": list(self.bus.motors.keys()),
|
||||
"shoulder": ["shoulder_pitch", "shoulder_yaw", "shoulder_roll"],
|
||||
@@ -122,11 +167,17 @@ class HopeJrArm(Robot):
|
||||
print("Calibration saved to", self.calibration_fpath)
|
||||
|
||||
def configure(self) -> None:
|
||||
"""Apply the operating mode, gains and limits from the configuration to the robot."""
|
||||
with self.bus.torque_disabled():
|
||||
self.bus.configure_motors(maximum_acceleration=30, acceleration=30)
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
# TODO: add docstring
|
||||
"""Assign each motor its bus ID, one at a time.
|
||||
|
||||
Run this once when building the robot. Interactive: prompts you to connect the controller board to a
|
||||
single motor at a time.
|
||||
"""
|
||||
for motor in reversed(self.bus.motors):
|
||||
input(f"Connect the controller board to the '{motor}' motor only and press enter.")
|
||||
self.bus.setup_motor(motor)
|
||||
@@ -135,6 +186,14 @@ class HopeJrArm(Robot):
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
# Read arm position
|
||||
"""Read the robot's current state and a frame from each camera.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: Keys matching [`~robots.Robot.observation_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
start = time.perf_counter()
|
||||
obs_dict = self.bus.sync_read("Present_Position", self.other_motors)
|
||||
obs_dict[self.shoulder_pitch] = self.bus.read("Present_Position", self.shoulder_pitch)
|
||||
@@ -160,6 +219,18 @@ class HopeJrArm(Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def send_action(self, action: RobotAction) -> RobotAction:
|
||||
"""Command the robot to move towards a target configuration.
|
||||
|
||||
Args:
|
||||
action (`dict[str, Any]`):
|
||||
Target values, keyed as in [`~robots.Robot.action_features`].
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The action actually sent, which may be clipped by `max_relative_target`.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
goal_pos = {key.removesuffix(".pos"): val for key, val in action.items() if key.endswith(".pos")}
|
||||
|
||||
# Cap goal position when too far away from present position.
|
||||
@@ -174,6 +245,11 @@ class HopeJrArm(Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self):
|
||||
"""Disconnect from the robot and its cameras.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
self.bus.disconnect(self.config.disable_torque_on_disconnect)
|
||||
for cam in self.cameras.values():
|
||||
cam.disconnect()
|
||||
|
||||
@@ -59,10 +59,22 @@ LEFT_HAND_INVERSIONS = [
|
||||
|
||||
|
||||
class HopeJrHand(Robot):
|
||||
"""One hand of the Hope Jr humanoid.
|
||||
|
||||
Each hand is its own robot, so a two-handed setup uses two of these with different `side` values. See
|
||||
[`~robots.Robot`] for the contract every method here implements.
|
||||
"""
|
||||
|
||||
config_class = HopeJrHandConfig
|
||||
name = "hope_jr_hand"
|
||||
|
||||
def __init__(self, config: HopeJrHandConfig):
|
||||
"""Build the robot from its configuration.
|
||||
|
||||
Args:
|
||||
config (`HopeJrHandConfig`):
|
||||
The robot's configuration. Its `port` and `cameras` determine what is connected.
|
||||
"""
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
self.bus = FeetechMotorsBus(
|
||||
@@ -113,18 +125,43 @@ class HopeJrHand(Robot):
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
"""The values this robot reports, and their types or shapes.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys as returned by [`~robots.Robot.get_observation`], mapped to a scalar type for
|
||||
proprioceptive values or to a `(height, width, channels)` shape for images.
|
||||
"""
|
||||
return {**self._motors_ft, **self._cameras_ft}
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
"""The values this robot accepts, and their types.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys accepted by [`~robots.Robot.send_action`], mapped to their type.
|
||||
"""
|
||||
return self._motors_ft
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Whether every device this robot uses is connected.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` only when the robot and all its cameras are connected.
|
||||
"""
|
||||
return self.bus.is_connected and all(cam.is_connected for cam in self.cameras.values())
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""Connect to the robot and its cameras, then apply the configured settings.
|
||||
|
||||
Args:
|
||||
calibrate (`bool`, *optional*, defaults to `True`):
|
||||
Whether to run calibration if the robot is not already calibrated.
|
||||
|
||||
Raises:
|
||||
DeviceAlreadyConnectedError: If the robot is already connected.
|
||||
"""
|
||||
self.bus.connect()
|
||||
if not self.is_calibrated and calibrate:
|
||||
self.calibrate()
|
||||
@@ -138,9 +175,18 @@ class HopeJrHand(Robot):
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
"""Whether the robot is calibrated.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` when no calibration is needed before use.
|
||||
"""
|
||||
return self.bus.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
"""Calibrate the robot and store the result.
|
||||
|
||||
Interactive: prompts on stdin and asks you to move the robot through the required positions.
|
||||
"""
|
||||
fingers = {}
|
||||
for finger in ["thumb", "index", "middle", "ring", "pinky"]:
|
||||
fingers[finger] = [motor for motor in self.bus.motors if motor.startswith(finger)]
|
||||
@@ -152,11 +198,17 @@ class HopeJrHand(Robot):
|
||||
print("Calibration saved to", self.calibration_fpath)
|
||||
|
||||
def configure(self) -> None:
|
||||
"""Apply the operating mode, gains and limits from the configuration to the robot."""
|
||||
with self.bus.torque_disabled():
|
||||
self.bus.configure_motors()
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
# TODO: add docstring
|
||||
"""Assign each motor its bus ID, one at a time.
|
||||
|
||||
Run this once when building the robot. Interactive: prompts you to connect the controller board to a
|
||||
single motor at a time.
|
||||
"""
|
||||
for motor in self.bus.motors:
|
||||
input(f"Connect the controller board to the '{motor}' motor only and press enter.")
|
||||
self.bus.setup_motor(motor)
|
||||
@@ -164,6 +216,14 @@ class HopeJrHand(Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
"""Read the robot's current state and a frame from each camera.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: Keys matching [`~robots.Robot.observation_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
obs_dict = {}
|
||||
|
||||
# Read hand position
|
||||
@@ -191,12 +251,29 @@ class HopeJrHand(Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def send_action(self, action: RobotAction) -> RobotAction:
|
||||
"""Command the robot to move towards a target configuration.
|
||||
|
||||
Args:
|
||||
action (`dict[str, Any]`):
|
||||
Target values, keyed as in [`~robots.Robot.action_features`].
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The action actually sent, which may be clipped by `max_relative_target`.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
goal_pos = {key.removesuffix(".pos"): val for key, val in action.items() if key.endswith(".pos")}
|
||||
self.bus.sync_write("Goal_Position", goal_pos)
|
||||
return action
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self):
|
||||
"""Disconnect from the robot and its cameras.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
self.bus.disconnect(self.config.disable_torque_on_disconnect)
|
||||
for cam in self.cameras.values():
|
||||
cam.disconnect()
|
||||
|
||||
@@ -22,6 +22,30 @@ from ..config import RobotConfig
|
||||
@RobotConfig.register_subclass("koch_follower")
|
||||
@dataclass
|
||||
class KochFollowerConfig(RobotConfig):
|
||||
"""Configuration for the Koch v1.1 follower arm.
|
||||
|
||||
Args:
|
||||
port (`str`):
|
||||
Serial port the arm is connected to, e.g. `/dev/ttyACM0` on Linux or `COM3` on Windows. Run
|
||||
`lerobot-find-port` to identify it.
|
||||
disable_torque_on_disconnect (`bool`, *optional*, defaults to `True`):
|
||||
Whether to release the motors on disconnect. Leave `True` unless the arm is holding a load it
|
||||
must not drop.
|
||||
max_relative_target (`float | dict[str, float]`, *optional*):
|
||||
Caps how far a single action may move the arm from its present position, as a safety limit. A
|
||||
scalar applies to every motor; a dict maps motor name to a per-motor cap. `None` disables
|
||||
clipping.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
Cameras to read alongside the arm's joint positions, keyed by the name they appear under in
|
||||
observations. Each must specify `width`, `height` and `fps`.
|
||||
use_degrees (`bool`, *optional*, defaults to `False`):
|
||||
Whether to report and accept joint positions in degrees rather than as a normalised range.
|
||||
id (`str`, *optional*):
|
||||
Identifier for this particular arm; also names its calibration file.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
|
||||
"""
|
||||
|
||||
# Port to connect to the arm
|
||||
port: str
|
||||
|
||||
|
||||
@@ -35,16 +35,24 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class KochFollower(Robot):
|
||||
"""
|
||||
- [Koch v1.0](https://github.com/AlexanderKoch-Koch/low_cost_robot), with and without the wrist-to-elbow
|
||||
expansion, developed by Alexander Koch from [Tau Robotics](https://tau-robotics.com)
|
||||
- [Koch v1.1](https://github.com/jess-moss/koch-v1-1) developed by Jess Moss
|
||||
"""The Koch follower arm, in either of its two revisions.
|
||||
|
||||
- [Koch v1.0](https://github.com/AlexanderKoch-Koch/low_cost_robot), with and without the
|
||||
wrist-to-elbow expansion, developed by Alexander Koch from
|
||||
[Tau Robotics](https://tau-robotics.com).
|
||||
- [Koch v1.1](https://github.com/jess-moss/koch-v1-1), developed by Jess Moss.
|
||||
"""
|
||||
|
||||
config_class = KochFollowerConfig
|
||||
name = "koch_follower"
|
||||
|
||||
def __init__(self, config: KochFollowerConfig):
|
||||
"""Build the robot from its configuration.
|
||||
|
||||
Args:
|
||||
config (`KochFollowerConfig`):
|
||||
The robot's configuration. Its `port` and `cameras` determine what is connected.
|
||||
"""
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
norm_mode_body = MotorNormMode.DEGREES if config.use_degrees else MotorNormMode.RANGE_M100_100
|
||||
@@ -79,23 +87,47 @@ class KochFollower(Robot):
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
"""The values this robot reports, and their types or shapes.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys as returned by [`~robots.Robot.get_observation`], mapped to a scalar type for
|
||||
proprioceptive values or to a `(height, width, channels)` shape for images.
|
||||
"""
|
||||
return {**self._motors_ft, **self._cameras_ft}
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
"""The values this robot accepts, and their types.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys accepted by [`~robots.Robot.send_action`], mapped to their type.
|
||||
"""
|
||||
return self._motors_ft
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Whether every device this robot uses is connected.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` only when the robot and all its cameras are connected.
|
||||
"""
|
||||
return self.bus.is_connected and all(cam.is_connected for cam in self.cameras.values())
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""
|
||||
We assume that at connection time, arm is in a rest position,
|
||||
and torque can be safely disabled to run calibration.
|
||||
"""
|
||||
"""Connect the motor bus and cameras, calibrating and configuring the arm.
|
||||
|
||||
> [!WARNING]
|
||||
> The arm is assumed to be at rest when this is called, because torque is disabled to run
|
||||
> calibration.
|
||||
|
||||
Args:
|
||||
calibrate (`bool`, *optional*, defaults to `True`):
|
||||
Whether to run calibration if the arm is not already calibrated.
|
||||
|
||||
Raises:
|
||||
DeviceAlreadyConnectedError: If the robot is already connected.
|
||||
"""
|
||||
self.bus.connect()
|
||||
if not self.is_calibrated and calibrate:
|
||||
logger.info(
|
||||
@@ -111,9 +143,18 @@ class KochFollower(Robot):
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
"""Whether the robot is calibrated.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` when no calibration is needed before use.
|
||||
"""
|
||||
return self.bus.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
"""Calibrate the robot and store the result.
|
||||
|
||||
Interactive: prompts on stdin and asks you to move the robot through the required positions.
|
||||
"""
|
||||
self.bus.disable_torque()
|
||||
if self.calibration:
|
||||
# Calibration file exists, ask user whether to use it or run new calibration
|
||||
@@ -157,6 +198,7 @@ class KochFollower(Robot):
|
||||
logger.info(f"Calibration saved to {self.calibration_fpath}")
|
||||
|
||||
def configure(self) -> None:
|
||||
"""Apply the operating mode, gains and limits from the configuration to the robot."""
|
||||
with self.bus.torque_disabled():
|
||||
self.bus.configure_motors()
|
||||
# Use 'extended position mode' for all motors except gripper, because in joint mode the servos
|
||||
@@ -181,6 +223,11 @@ class KochFollower(Robot):
|
||||
self.bus.write("Position_D_Gain", "elbow_flex", 600)
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
"""Assign each motor its bus ID, one at a time.
|
||||
|
||||
Run this once when building the robot. Interactive: prompts you to connect the controller board to a
|
||||
single motor at a time.
|
||||
"""
|
||||
for motor in reversed(self.bus.motors):
|
||||
input(f"Connect the controller board to the '{motor}' motor only and press enter.")
|
||||
self.bus.setup_motor(motor)
|
||||
@@ -189,6 +236,14 @@ class KochFollower(Robot):
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
# Read arm position
|
||||
"""Read the robot's current state and a frame from each camera.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: Keys matching [`~robots.Robot.observation_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
start = time.perf_counter()
|
||||
obs_dict = self.bus.sync_read("Present_Position")
|
||||
obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
|
||||
@@ -225,7 +280,6 @@ class KochFollower(Robot):
|
||||
Returns:
|
||||
RobotAction: The action sent to the motors, potentially clipped.
|
||||
"""
|
||||
|
||||
goal_pos = {key.removesuffix(".pos"): val for key, val in action.items() if key.endswith(".pos")}
|
||||
|
||||
# Cap goal position when too far away from present position.
|
||||
@@ -241,6 +295,11 @@ class KochFollower(Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self):
|
||||
"""Disconnect from the robot and its cameras.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
self.bus.disconnect(self.config.disable_torque_on_disconnect)
|
||||
for cam in self.cameras.values():
|
||||
cam.disconnect()
|
||||
|
||||
@@ -21,6 +21,12 @@ from ..config import RobotConfig
|
||||
|
||||
|
||||
def lekiwi_cameras_config() -> dict[str, CameraConfig]:
|
||||
"""Build the default camera set for a LeKiwi base.
|
||||
|
||||
Returns:
|
||||
`dict[str, CameraConfig]`: The `front` and `wrist` OpenCV cameras at the device paths and
|
||||
rotations of a standard LeKiwi build. Override the `cameras` field if yours is wired differently.
|
||||
"""
|
||||
return {
|
||||
"front": OpenCVCameraConfig(
|
||||
index_or_path="/dev/video0",
|
||||
@@ -44,6 +50,34 @@ def lekiwi_cameras_config() -> dict[str, CameraConfig]:
|
||||
@RobotConfig.register_subclass("lekiwi")
|
||||
@dataclass
|
||||
class LeKiwiConfig(RobotConfig):
|
||||
"""Configuration for LeKiwi, running on the robot itself.
|
||||
|
||||
This is the config used by the process on the LeKiwi's own computer. To drive one from another machine,
|
||||
use [`LeKiwiClientConfig`] instead.
|
||||
|
||||
Args:
|
||||
port (`str`, *optional*, defaults to `"/dev/ttyACM0"`):
|
||||
Serial port of the motor bus on the robot's computer.
|
||||
disable_torque_on_disconnect (`bool`, *optional*, defaults to `True`):
|
||||
Whether to release the motors on disconnect.
|
||||
max_relative_target (`float | dict[str, float]`, *optional*):
|
||||
Caps how far a single action may move the arm from its present position, as a safety limit. A
|
||||
scalar applies to every motor; a dict maps motor name to a per-motor cap. `None` disables
|
||||
clipping.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
Cameras to read alongside the joint positions. Defaults to the standard `front` and `wrist`
|
||||
build; see [`lekiwi_cameras_config`].
|
||||
use_degrees (`bool`, *optional*, defaults to `True`):
|
||||
Whether to report and accept arm joint positions in degrees.
|
||||
num_read_retries (`int`, *optional*, defaults to 2):
|
||||
Extra attempts when a `sync_read` fails. Feetech buses occasionally return a corrupted status
|
||||
packet, which would otherwise abort the control loop.
|
||||
id (`str`, *optional*):
|
||||
Identifier for this particular robot; also names its calibration file.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
|
||||
"""
|
||||
|
||||
port: str = "/dev/ttyACM0" # port to connect to the bus
|
||||
|
||||
disable_torque_on_disconnect: bool = True
|
||||
@@ -67,6 +101,22 @@ class LeKiwiConfig(RobotConfig):
|
||||
|
||||
@dataclass
|
||||
class LeKiwiHostConfig:
|
||||
"""Configuration for the host process that serves a LeKiwi over the network.
|
||||
|
||||
Args:
|
||||
port_zmq_cmd (`int`, *optional*, defaults to 5555):
|
||||
ZMQ port the host listens on for actions.
|
||||
port_zmq_observations (`int`, *optional*, defaults to 5556):
|
||||
ZMQ port the host publishes observations on.
|
||||
connection_time_s (`int`, *optional*, defaults to 30):
|
||||
How long the host stays up before shutting down.
|
||||
watchdog_timeout_ms (`int`, *optional*, defaults to 500):
|
||||
Stop the robot if no command arrives within this window. Guards against a dropped client
|
||||
leaving the base driving.
|
||||
max_loop_freq_hz (`int`, *optional*, defaults to 30):
|
||||
Control loop frequency. Lower it if the robot jitters, and watch CPU load with `top`.
|
||||
"""
|
||||
|
||||
# Network Configuration
|
||||
port_zmq_cmd: int = 5555
|
||||
port_zmq_observations: int = 5556
|
||||
@@ -84,6 +134,32 @@ class LeKiwiHostConfig:
|
||||
@RobotConfig.register_subclass("lekiwi_client")
|
||||
@dataclass
|
||||
class LeKiwiClientConfig(RobotConfig):
|
||||
"""Configuration for driving a LeKiwi from another machine.
|
||||
|
||||
Presents the same [`~robots.Robot`] interface as the robot-side [`LeKiwiConfig`], but every call goes
|
||||
over ZMQ to the host process. Calibration lives on the robot, so nothing here configures it.
|
||||
|
||||
Args:
|
||||
remote_ip (`str`):
|
||||
IP address of the LeKiwi's computer on the network.
|
||||
port_zmq_cmd (`int`, *optional*, defaults to 5555):
|
||||
ZMQ port to send actions to. Must match the host's `port_zmq_cmd`.
|
||||
port_zmq_observations (`int`, *optional*, defaults to 5556):
|
||||
ZMQ port to receive observations on. Must match the host's `port_zmq_observations`.
|
||||
teleop_keys (`dict[str, str]`, *optional*):
|
||||
Keyboard bindings for driving the base: movement, rotation, speed control and quit.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
Cameras expected in the observation stream. Defaults to the standard `front` and `wrist` build.
|
||||
polling_timeout_ms (`int`, *optional*, defaults to 15):
|
||||
How long to wait for an observation before giving up on that step.
|
||||
connect_timeout_s (`int`, *optional*, defaults to 5):
|
||||
How long to wait for the host to answer when connecting.
|
||||
id (`str`, *optional*):
|
||||
Identifier for this particular robot.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Unused by the client: calibration is held on the robot.
|
||||
"""
|
||||
|
||||
# Network Configuration
|
||||
remote_ip: str
|
||||
port_zmq_cmd: int = 5555
|
||||
|
||||
@@ -39,17 +39,25 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LeKiwi(Robot):
|
||||
"""
|
||||
The robot includes a three omniwheel mobile base and a remote follower arm.
|
||||
The leader arm is connected locally (on the laptop) and its joint positions are recorded and then
|
||||
forwarded to the remote follower arm (after applying a safety clamp).
|
||||
In parallel, keyboard teleoperation is used to generate raw velocity commands for the wheels.
|
||||
"""A three-omniwheel mobile base with a follower arm on top, running on the robot itself.
|
||||
|
||||
The leader arm is connected to the operator's laptop; its joint positions are recorded and forwarded
|
||||
to this follower arm after a safety clamp. In parallel, keyboard teleoperation generates raw velocity
|
||||
commands for the wheels.
|
||||
|
||||
To drive one of these from another machine, use [`~robots.lekiwi.LeKiwiClient`].
|
||||
"""
|
||||
|
||||
config_class = LeKiwiConfig
|
||||
name = "lekiwi"
|
||||
|
||||
def __init__(self, config: LeKiwiConfig):
|
||||
"""Build the robot from its configuration.
|
||||
|
||||
Args:
|
||||
config (`LeKiwiConfig`):
|
||||
The robot's configuration. Its `port` and `cameras` determine what is connected.
|
||||
"""
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
norm_mode_body = MotorNormMode.DEGREES if config.use_degrees else MotorNormMode.RANGE_M100_100
|
||||
@@ -105,18 +113,43 @@ class LeKiwi(Robot):
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
"""The values this robot reports, and their types or shapes.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys as returned by [`~robots.Robot.get_observation`], mapped to a scalar type for
|
||||
proprioceptive values or to a `(height, width, channels)` shape for images.
|
||||
"""
|
||||
return {**self._state_ft, **self._cameras_ft}
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
"""The values this robot accepts, and their types.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys accepted by [`~robots.Robot.send_action`], mapped to their type.
|
||||
"""
|
||||
return self._state_ft
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Whether every device this robot uses is connected.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` only when the robot and all its cameras are connected.
|
||||
"""
|
||||
return self.bus.is_connected and all(cam.is_connected for cam in self.cameras.values())
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""Connect to the robot and its cameras, then apply the configured settings.
|
||||
|
||||
Args:
|
||||
calibrate (`bool`, *optional*, defaults to `True`):
|
||||
Whether to run calibration if the robot is not already calibrated.
|
||||
|
||||
Raises:
|
||||
DeviceAlreadyConnectedError: If the robot is already connected.
|
||||
"""
|
||||
self.bus.connect()
|
||||
if not self.is_calibrated and calibrate:
|
||||
logger.info(
|
||||
@@ -132,9 +165,18 @@ class LeKiwi(Robot):
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
"""Whether the robot is calibrated.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` when no calibration is needed before use.
|
||||
"""
|
||||
return self.bus.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
"""Calibrate the robot and store the result.
|
||||
|
||||
Interactive: prompts on stdin and asks you to move the robot through the required positions.
|
||||
"""
|
||||
if self.calibration:
|
||||
# Calibration file exists, ask user whether to use it or run new calibration
|
||||
user_input = input(
|
||||
@@ -189,6 +231,7 @@ class LeKiwi(Robot):
|
||||
# Set-up arm actuators (position mode)
|
||||
# We assume that at connection time, arm is in a rest position,
|
||||
# and torque can be safely disabled to run calibration.
|
||||
"""Apply the operating mode, gains and limits from the configuration to the robot."""
|
||||
self.bus.disable_torque()
|
||||
self.bus.configure_motors()
|
||||
for name in self.arm_motors:
|
||||
@@ -205,6 +248,11 @@ class LeKiwi(Robot):
|
||||
self.bus.enable_torque()
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
"""Assign each motor its bus ID, one at a time.
|
||||
|
||||
Run this once when building the robot. Interactive: prompts you to connect the controller board to a
|
||||
single motor at a time.
|
||||
"""
|
||||
for motor in chain(reversed(self.arm_motors), reversed(self.base_motors)):
|
||||
input(f"Connect the controller board to the '{motor}' motor only and press enter.")
|
||||
self.bus.setup_motor(motor)
|
||||
@@ -238,8 +286,7 @@ class LeKiwi(Robot):
|
||||
base_radius: float = 0.125,
|
||||
max_raw: int = 3000,
|
||||
) -> dict:
|
||||
"""
|
||||
Convert desired body-frame velocities into wheel raw commands.
|
||||
"""Convert desired body-frame velocities into wheel raw commands.
|
||||
|
||||
Parameters:
|
||||
x_cmd : Linear velocity in x (m/s).
|
||||
@@ -302,8 +349,7 @@ class LeKiwi(Robot):
|
||||
wheel_radius: float = 0.05,
|
||||
base_radius: float = 0.125,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Convert wheel raw command feedback back into body-frame velocities.
|
||||
"""Convert wheel raw command feedback back into body-frame velocities.
|
||||
|
||||
Parameters:
|
||||
wheel_raw : Vector with raw wheel commands ("base_left_wheel", "base_back_wheel", "base_right_wheel").
|
||||
@@ -313,7 +359,6 @@ class LeKiwi(Robot):
|
||||
Returns:
|
||||
A dict (x.vel, y.vel, theta.vel) all in m/s
|
||||
"""
|
||||
|
||||
# Convert each raw command back to an angular speed in deg/s.
|
||||
wheel_degps = np.array(
|
||||
[
|
||||
@@ -346,6 +391,14 @@ class LeKiwi(Robot):
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
# Read actuators position for arm and vel for base
|
||||
"""Read the robot's current state and a frame from each camera.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: Keys matching [`~robots.Robot.observation_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
start = time.perf_counter()
|
||||
arm_pos = self.bus.sync_read(
|
||||
"Present_Position", self.arm_motors, num_retry=self.config.num_read_retries
|
||||
@@ -390,7 +443,6 @@ class LeKiwi(Robot):
|
||||
Returns:
|
||||
RobotAction: the action sent to the motors, potentially clipped.
|
||||
"""
|
||||
|
||||
arm_goal_pos = {k: v for k, v in action.items() if k.endswith(".pos")}
|
||||
base_goal_vel = {k: v for k, v in action.items() if k.endswith(".vel")}
|
||||
|
||||
@@ -419,11 +471,17 @@ class LeKiwi(Robot):
|
||||
return {**arm_goal_pos, **base_goal_vel}
|
||||
|
||||
def stop_base(self):
|
||||
"""Bring the mobile base to a halt by commanding zero velocity on its wheels."""
|
||||
self.bus.sync_write("Goal_Velocity", dict.fromkeys(self.base_motors, 0), num_retry=5)
|
||||
logger.info("Base motors stopped")
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self):
|
||||
"""Disconnect from the robot and its cameras.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
self.stop_base()
|
||||
self.bus.disconnect(self.config.disable_torque_on_disconnect)
|
||||
for cam in self.cameras.values():
|
||||
|
||||
@@ -31,10 +31,23 @@ from .config_lekiwi import LeKiwiClientConfig
|
||||
|
||||
|
||||
class LeKiwiClient(Robot):
|
||||
"""Drives a LeKiwi over the network from another machine.
|
||||
|
||||
Presents the same [`~robots.Robot`] interface as [`~robots.lekiwi.LeKiwi`], but every observation and
|
||||
action crosses a ZMQ connection to the host process running on the robot. Calibration stays on the
|
||||
robot, so this class does not perform it.
|
||||
"""
|
||||
|
||||
config_class = LeKiwiClientConfig
|
||||
name = "lekiwi_client"
|
||||
|
||||
def __init__(self, config: LeKiwiClientConfig):
|
||||
"""Build the robot from its configuration.
|
||||
|
||||
Args:
|
||||
config (`LeKiwiClientConfig`):
|
||||
The robot's configuration. Its `port` and `cameras` determine what is connected.
|
||||
"""
|
||||
import zmq
|
||||
|
||||
self._zmq = zmq
|
||||
@@ -105,24 +118,50 @@ class LeKiwiClient(Robot):
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
"""The values this robot reports, and their types or shapes.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys as returned by [`~robots.Robot.get_observation`], mapped to a scalar type for
|
||||
proprioceptive values or to a `(height, width, channels)` shape for images.
|
||||
"""
|
||||
return {**self._state_ft, **self._cameras_ft}
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
"""The values this robot accepts, and their types.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys accepted by [`~robots.Robot.send_action`], mapped to their type.
|
||||
"""
|
||||
return self._state_ft
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Whether every device this robot uses is connected.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` only when the robot and all its cameras are connected.
|
||||
"""
|
||||
return self._is_connected
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
"""Whether the robot is calibrated.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` when no calibration is needed before use.
|
||||
"""
|
||||
pass
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self) -> None:
|
||||
"""Establishes ZMQ sockets with the remote mobile robot"""
|
||||
"""Open the ZMQ command and observation sockets to the LeKiwi host.
|
||||
|
||||
Takes no `calibrate` argument: calibration lives on the robot, not on the client.
|
||||
|
||||
Raises:
|
||||
DeviceAlreadyConnectedError: If the client is already connected.
|
||||
"""
|
||||
zmq = self._zmq
|
||||
self.zmq_context = zmq.Context()
|
||||
self.zmq_cmd_socket = self.zmq_context.socket(zmq.PUSH)
|
||||
@@ -146,6 +185,10 @@ class LeKiwiClient(Robot):
|
||||
self._is_connected = True
|
||||
|
||||
def calibrate(self) -> None:
|
||||
"""Calibrate the robot and store the result.
|
||||
|
||||
Interactive: prompts on stdin and asks you to move the robot through the required positions.
|
||||
"""
|
||||
pass
|
||||
|
||||
def _poll_and_get_latest_message(self) -> list[bytes] | None:
|
||||
@@ -203,7 +246,6 @@ class LeKiwiClient(Robot):
|
||||
self, observation: RobotObservation
|
||||
) -> tuple[dict[str, np.ndarray], RobotObservation]:
|
||||
"""Extracts frames, and state from the parsed observation."""
|
||||
|
||||
flat_state = {key: observation.get(key, 0.0) for key in self._state_order}
|
||||
|
||||
state_vec = np.array([flat_state[key] for key in self._state_order], dtype=np.float32)
|
||||
@@ -222,14 +264,12 @@ class LeKiwiClient(Robot):
|
||||
return current_frames, obs_dict
|
||||
|
||||
def _get_data(self) -> tuple[dict[str, np.ndarray], RobotObservation]:
|
||||
"""
|
||||
Polls the video socket for the latest observation data.
|
||||
"""Polls the video socket for the latest observation data.
|
||||
|
||||
Attempts to retrieve and decode the latest message within a short timeout.
|
||||
If successful, updates and returns the new frames, speed, and arm state.
|
||||
If no new data arrives or decoding fails, returns the last known values.
|
||||
"""
|
||||
|
||||
# 1. Get the latest message's frames from the socket
|
||||
latest_frames = self._poll_and_get_latest_message()
|
||||
|
||||
@@ -258,12 +298,17 @@ class LeKiwiClient(Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
"""
|
||||
Capture observations from the remote robot: current follower arm positions,
|
||||
present wheel speeds (converted to body-frame velocities: x, y, theta),
|
||||
and a camera frame. Receives over ZMQ, translate to body-frame vel
|
||||
"""
|
||||
"""Receive one observation from the remote robot over ZMQ.
|
||||
|
||||
Wheel speeds arrive as raw motor velocities and are converted here to body-frame `x`, `y` and
|
||||
`theta`.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: Follower arm positions, body-frame base velocities and camera frames.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the client is not connected.
|
||||
"""
|
||||
frames, obs_dict = self._get_data()
|
||||
|
||||
# Loop over each configured camera
|
||||
@@ -308,21 +353,24 @@ class LeKiwiClient(Robot):
|
||||
}
|
||||
|
||||
def configure(self):
|
||||
"""Apply the operating mode, gains and limits from the configuration to the robot."""
|
||||
pass
|
||||
|
||||
@check_if_not_connected
|
||||
def send_action(self, action: RobotAction) -> RobotAction:
|
||||
"""Command lekiwi to move to a target joint configuration. Translates to motor space + sends over ZMQ
|
||||
"""Send a target configuration to the remote robot over ZMQ.
|
||||
|
||||
Body-frame base velocities are translated into wheel velocities before sending.
|
||||
|
||||
Args:
|
||||
action (RobotAction): array containing the goal positions for the motors.
|
||||
action (`dict[str, Any]`): Goal positions for the arm and body-frame velocities for the base.
|
||||
|
||||
Raises:
|
||||
RobotDeviceNotConnectedError: if robot is not connected.
|
||||
|
||||
Returns:
|
||||
np.ndarray: the action sent to the motors, potentially clipped.
|
||||
"""
|
||||
|
||||
# Action values may be torch tensors (e.g. replayed from a dataset) or numpy
|
||||
# scalars; json.dumps only serializes Python primitives, so coerce each value to a
|
||||
# plain float before sending.
|
||||
@@ -338,8 +386,11 @@ class LeKiwiClient(Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self):
|
||||
"""Cleans ZMQ comms"""
|
||||
"""Close the ZMQ sockets and terminate the context.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the client is not connected.
|
||||
"""
|
||||
self.zmq_observation_socket.close()
|
||||
self.zmq_cmd_socket.close()
|
||||
self.zmq_context.term()
|
||||
|
||||
@@ -36,7 +36,19 @@ class LeKiwiServerConfig:
|
||||
|
||||
|
||||
class LeKiwiHost:
|
||||
"""Serves a [`~robots.lekiwi.LeKiwi`] over ZMQ so a client can drive it from another machine.
|
||||
|
||||
Runs on the robot's own computer, receiving actions on one socket and publishing observations on
|
||||
another.
|
||||
"""
|
||||
|
||||
def __init__(self, config: LeKiwiHostConfig):
|
||||
"""Bind the command and observation sockets.
|
||||
|
||||
Args:
|
||||
config (`LeKiwiHostConfig`):
|
||||
Ports, loop frequency and watchdog settings for the host.
|
||||
"""
|
||||
self.zmq_context = zmq.Context()
|
||||
self.zmq_cmd_socket = self.zmq_context.socket(zmq.PULL)
|
||||
self.zmq_cmd_socket.setsockopt(zmq.CONFLATE, 1)
|
||||
@@ -53,6 +65,11 @@ class LeKiwiHost:
|
||||
self.max_loop_freq_hz = config.max_loop_freq_hz
|
||||
|
||||
def disconnect(self):
|
||||
"""Disconnect from the robot and its cameras.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
self.zmq_observation_socket.close()
|
||||
self.zmq_cmd_socket.close()
|
||||
self.zmq_context.term()
|
||||
@@ -60,6 +77,12 @@ class LeKiwiHost:
|
||||
|
||||
@draccus.wrap()
|
||||
def main(cfg: LeKiwiServerConfig):
|
||||
"""Run the LeKiwi host loop until the configured connection time elapses.
|
||||
|
||||
Args:
|
||||
cfg (`LeKiwiServerConfig`):
|
||||
The robot and host configuration to serve.
|
||||
"""
|
||||
logging.info("Configuring LeKiwi")
|
||||
robot = LeKiwi(cfg.robot)
|
||||
|
||||
|
||||
@@ -22,6 +22,30 @@ from ..config import RobotConfig
|
||||
@RobotConfig.register_subclass("omx_follower")
|
||||
@dataclass
|
||||
class OmxFollowerConfig(RobotConfig):
|
||||
"""Configuration for the OpenMANIPULATOR-X follower arm.
|
||||
|
||||
Args:
|
||||
port (`str`):
|
||||
Serial port the arm is connected to, e.g. `/dev/ttyUSB0`. Run `lerobot-find-port` to identify
|
||||
it.
|
||||
disable_torque_on_disconnect (`bool`, *optional*, defaults to `True`):
|
||||
Whether to release the motors on disconnect. Leave `True` unless the arm is holding a load it
|
||||
must not drop.
|
||||
max_relative_target (`float | dict[str, float]`, *optional*):
|
||||
Caps how far a single action may move the arm from its present position, as a safety limit. A
|
||||
scalar applies to every motor; a dict maps motor name to a per-motor cap. `None` disables
|
||||
clipping.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
Cameras to read alongside the arm's joint positions, keyed by the name they appear under in
|
||||
observations. Each must specify `width`, `height` and `fps`.
|
||||
use_degrees (`bool`, *optional*, defaults to `False`):
|
||||
Whether to report and accept joint positions in degrees rather than as a normalised range.
|
||||
id (`str`, *optional*):
|
||||
Identifier for this particular arm; also names its calibration file.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
|
||||
"""
|
||||
|
||||
# Port to connect to the arm
|
||||
port: str
|
||||
|
||||
|
||||
@@ -36,15 +36,21 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OmxFollower(Robot):
|
||||
"""
|
||||
- [OMX](https://github.com/ROBOTIS-GIT/open_manipulator),
|
||||
expansion, developed by Woojin Wie and Junha Cha from [ROBOTIS](https://ai.robotis.com/)
|
||||
"""The [OpenMANIPULATOR-X](https://github.com/ROBOTIS-GIT/open_manipulator) follower arm.
|
||||
|
||||
Developed by Woojin Wie and Junha Cha at [ROBOTIS](https://ai.robotis.com/).
|
||||
"""
|
||||
|
||||
config_class = OmxFollowerConfig
|
||||
name = "omx_follower"
|
||||
|
||||
def __init__(self, config: OmxFollowerConfig):
|
||||
"""Build the robot from its configuration.
|
||||
|
||||
Args:
|
||||
config (`OmxFollowerConfig`):
|
||||
The robot's configuration. Its `port` and `cameras` determine what is connected.
|
||||
"""
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
norm_mode_body = MotorNormMode.DEGREES if config.use_degrees else MotorNormMode.RANGE_M100_100
|
||||
@@ -79,25 +85,50 @@ class OmxFollower(Robot):
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
"""The values this robot reports, and their types or shapes.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys as returned by [`~robots.Robot.get_observation`], mapped to a scalar type for
|
||||
proprioceptive values or to a `(height, width, channels)` shape for images.
|
||||
"""
|
||||
return {**self._motors_ft, **self._cameras_ft}
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
"""The values this robot accepts, and their types.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys accepted by [`~robots.Robot.send_action`], mapped to their type.
|
||||
"""
|
||||
return self._motors_ft
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Whether every device this robot uses is connected.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` only when the robot and all its cameras are connected.
|
||||
"""
|
||||
return self.bus.is_connected and all(cam.is_connected for cam in self.cameras.values())
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""
|
||||
For OMX robots that come pre-calibrated:
|
||||
- If default calibration from package doesn't match motors, read from motors and save
|
||||
- This allows using pre-calibrated robots without manual calibration
|
||||
- If no calibration file exists, use factory default values (homing_offset=0, range_min=0, range_max=4095)
|
||||
"""
|
||||
"""Connect the motor bus and cameras, handling the pre-calibrated case.
|
||||
|
||||
OMX arms ship calibrated, so this avoids asking for a manual calibration where possible:
|
||||
|
||||
- if the packaged default calibration does not match the motors, the motors' own values are read
|
||||
and saved;
|
||||
- if no calibration file exists, factory defaults are used (`homing_offset=0`, `range_min=0`,
|
||||
`range_max=4095`).
|
||||
|
||||
Args:
|
||||
calibrate (`bool`, *optional*, defaults to `True`):
|
||||
Whether to calibrate if the arm is not already calibrated.
|
||||
|
||||
Raises:
|
||||
DeviceAlreadyConnectedError: If the robot is already connected.
|
||||
"""
|
||||
self.bus.connect()
|
||||
if not self.is_calibrated and calibrate:
|
||||
logger.info(
|
||||
@@ -113,9 +144,18 @@ class OmxFollower(Robot):
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
"""Whether the robot is calibrated.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` when no calibration is needed before use.
|
||||
"""
|
||||
return self.bus.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
"""Calibrate the robot and store the result.
|
||||
|
||||
Interactive: prompts on stdin and asks you to move the robot through the required positions.
|
||||
"""
|
||||
self.bus.disable_torque()
|
||||
logger.info(f"\nUsing factory default calibration values for {self}")
|
||||
logger.info(f"\nWriting default configuration of {self} to the motors")
|
||||
@@ -140,6 +180,7 @@ class OmxFollower(Robot):
|
||||
logger.info(f"Calibration saved to {self.calibration_fpath}")
|
||||
|
||||
def configure(self) -> None:
|
||||
"""Apply the operating mode, gains and limits from the configuration to the robot."""
|
||||
with self.bus.torque_disabled():
|
||||
self.bus.configure_motors()
|
||||
# Use 'extended position mode' for all motors except gripper, because in joint mode the servos
|
||||
@@ -164,6 +205,11 @@ class OmxFollower(Robot):
|
||||
self.bus.write("Position_D_Gain", "elbow_flex", 600)
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
"""Assign each motor its bus ID, one at a time.
|
||||
|
||||
Run this once when building the robot. Interactive: prompts you to connect the controller board to a
|
||||
single motor at a time.
|
||||
"""
|
||||
for motor in reversed(self.bus.motors):
|
||||
input(f"Connect the controller board to the '{motor}' motor only and press enter.")
|
||||
self.bus.setup_motor(motor)
|
||||
@@ -172,6 +218,14 @@ class OmxFollower(Robot):
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
# Read arm position
|
||||
"""Read the robot's current state and a frame from each camera.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: Keys matching [`~robots.Robot.observation_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
start = time.perf_counter()
|
||||
obs_dict = self.bus.sync_read("Present_Position")
|
||||
obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
|
||||
@@ -208,7 +262,6 @@ class OmxFollower(Robot):
|
||||
Returns:
|
||||
RobotAction: The action sent to the motors, potentially clipped.
|
||||
"""
|
||||
|
||||
goal_pos = {key.removesuffix(".pos"): val for key, val in action.items() if key.endswith(".pos")}
|
||||
|
||||
# Cap goal position when too far away from present position.
|
||||
@@ -224,6 +277,11 @@ class OmxFollower(Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self):
|
||||
"""Disconnect from the robot and its cameras.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
self.bus.disconnect(self.config.disable_torque_on_disconnect)
|
||||
for cam in self.cameras.values():
|
||||
cam.disconnect()
|
||||
|
||||
@@ -45,7 +45,13 @@ RIGHT_DEFAULT_JOINTS_LIMITS: dict[str, tuple[float, float]] = {
|
||||
|
||||
@dataclass
|
||||
class OpenArmFollowerConfigBase:
|
||||
"""Base configuration for the OpenArms follower robot with Damiao motors."""
|
||||
"""Field definitions for the OpenArm follower, a 7-DOF arm plus gripper on Damiao CAN motors.
|
||||
|
||||
This class only carries the fields. The registered configuration users instantiate is
|
||||
[`OpenArmFollowerConfig`], which documents them all in one place — doc-builder renders only a class's
|
||||
own docstring, never its bases'. It is also used directly as the per-arm config of
|
||||
[`~robots.bi_openarm_follower.BiOpenArmFollowerConfig`].
|
||||
"""
|
||||
|
||||
# CAN interfaces - one per arm
|
||||
# arm CAN interface (e.g., "can1")
|
||||
@@ -123,4 +129,57 @@ class OpenArmFollowerConfigBase:
|
||||
@RobotConfig.register_subclass("openarm_follower")
|
||||
@dataclass
|
||||
class OpenArmFollowerConfig(RobotConfig, OpenArmFollowerConfigBase):
|
||||
"""Configuration for a single OpenArm follower arm.
|
||||
|
||||
OpenArm is a 7-DOF arm plus gripper on Damiao CAN motors, so `port` names a CAN interface rather than a
|
||||
serial device. Calibration follows the usual LeRobot flow and is stored per `id`.
|
||||
|
||||
> [!WARNING]
|
||||
> `joint_limits` defaults to a deliberately tiny range so an uncalibrated arm cannot swing. Set `side`
|
||||
> to `"left"` or `"right"` to get the real limits for that arm, or pass your own.
|
||||
|
||||
The per-joint lists — `position_kp`, `position_kd` — hold 8 values in motor order: `joint_1` through
|
||||
`joint_7`, then `gripper`.
|
||||
|
||||
Args:
|
||||
port (`str`):
|
||||
CAN interface the arm is on, e.g. `"can0"` on Linux.
|
||||
side (`str`, *optional*):
|
||||
Which arm this is, `"left"` or `"right"`. Selects that side's joint limits. Leaving it `None`
|
||||
keeps the small safety defaults.
|
||||
can_interface (`str`, *optional*, defaults to `"socketcan"`):
|
||||
CAN backend: `"socketcan"` on Linux, `"slcan"` for a serial adapter, or `"auto"` to detect.
|
||||
use_can_fd (`bool`, *optional*, defaults to `True`):
|
||||
Whether to use CAN FD. OpenArm uses it by default.
|
||||
can_bitrate (`int`, *optional*, defaults to 1000000):
|
||||
Nominal CAN bitrate, 1 Mbps.
|
||||
can_data_bitrate (`int`, *optional*, defaults to 5000000):
|
||||
CAN FD data bitrate, 5 Mbps. Only used when `use_can_fd` is `True`.
|
||||
disable_torque_on_disconnect (`bool`, *optional*, defaults to `True`):
|
||||
Whether to release the motors on disconnect. Leave `True` unless the arm is holding a load it
|
||||
must not drop.
|
||||
use_velocity_and_torque (`bool`, *optional*, defaults to `False`):
|
||||
Whether to expose `.vel` and `.torque` per motor in the observation features. Kept `False` by
|
||||
default for compatibility with the position-only `openarm_mini` teleoperator.
|
||||
max_relative_target (`float | dict[str, float]`, *optional*):
|
||||
Caps how far a single action may move the arm from its present position, as a safety limit. A
|
||||
scalar applies to every motor; a dict maps motor name to a per-motor cap. `None` disables
|
||||
clipping.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
Cameras to read alongside the arm's joint positions.
|
||||
motor_config (`dict[str, tuple[int, int, str]]`, *optional*):
|
||||
Maps motor name to `(send_can_id, recv_can_id, motor_type)`. Defaults to the stock OpenArm
|
||||
layout; change it only if you have rewired or re-addressed the motors.
|
||||
position_kp (`list[float]`, *optional*):
|
||||
MIT-mode proportional gains used by `send_action`, 8 values in motor order.
|
||||
position_kd (`list[float]`, *optional*):
|
||||
MIT-mode derivative gains used by `send_action`, 8 values in motor order.
|
||||
joint_limits (`dict[str, tuple[float, float]]`, *optional*):
|
||||
Soft `(min, max)` limits in degrees per joint, clipped against on every action.
|
||||
id (`str`, *optional*):
|
||||
Identifier for this particular arm; also names its calibration file.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
@@ -37,15 +37,22 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OpenArmFollower(Robot):
|
||||
"""
|
||||
OpenArms Follower Robot which uses CAN bus communication to control 7 DOF arm with a gripper.
|
||||
The arm uses Damiao motors in MIT control mode.
|
||||
"""The OpenArm follower: a 7-DOF arm plus gripper on a CAN bus.
|
||||
|
||||
Uses Damiao motors in MIT control mode. See [`~robots.Robot`] for the contract every method here
|
||||
implements.
|
||||
"""
|
||||
|
||||
config_class = OpenArmFollowerConfig
|
||||
name = "openarm_follower"
|
||||
|
||||
def __init__(self, config: OpenArmFollowerConfig):
|
||||
"""Build the robot from its configuration.
|
||||
|
||||
Args:
|
||||
config (`OpenArmFollowerConfig`):
|
||||
The robot's configuration. Its `port` and `cameras` determine what is connected.
|
||||
"""
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
|
||||
@@ -127,13 +134,11 @@ class OpenArmFollower(Robot):
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""
|
||||
Connect to the robot and optionally calibrate.
|
||||
"""Connect to the robot and optionally calibrate.
|
||||
|
||||
We assume that at connection time, the arms are in a safe rest position,
|
||||
and torque can be safely disabled to run calibration if needed.
|
||||
"""
|
||||
|
||||
# Connect to CAN bus
|
||||
logger.info(f"Connecting arm on {self.config.port}...")
|
||||
self.bus.connect()
|
||||
@@ -160,8 +165,7 @@ class OpenArmFollower(Robot):
|
||||
return self.bus.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
"""
|
||||
Run calibration procedure for OpenArms robot.
|
||||
"""Run calibration procedure for OpenArms robot.
|
||||
|
||||
The calibration procedure:
|
||||
1. Disable torque
|
||||
@@ -217,14 +221,18 @@ class OpenArmFollower(Robot):
|
||||
self.bus.configure_motors()
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
"""Assign each motor its bus ID, one at a time.
|
||||
|
||||
Run this once when building the robot. Interactive: prompts you to connect the controller board to a
|
||||
single motor at a time.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Motor ID configuration is typically done via manufacturer tools for CAN motors."
|
||||
)
|
||||
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
"""
|
||||
Get current observation from robot including position, velocity, and torque.
|
||||
"""Get current observation from robot including position, velocity, and torque.
|
||||
|
||||
Reads all motor states (pos/vel/torque) in one CAN refresh cycle
|
||||
instead of 3 separate reads.
|
||||
@@ -268,8 +276,7 @@ class OpenArmFollower(Robot):
|
||||
custom_kp: dict[str, float] | None = None,
|
||||
custom_kd: dict[str, float] | None = None,
|
||||
) -> RobotAction:
|
||||
"""
|
||||
Send action command to robot.
|
||||
"""Send action command to robot.
|
||||
|
||||
The action magnitude may be clipped based on safety limits.
|
||||
|
||||
@@ -281,7 +288,6 @@ class OpenArmFollower(Robot):
|
||||
Returns:
|
||||
The action actually sent (potentially clipped)
|
||||
"""
|
||||
|
||||
goal_pos = {key.removesuffix(".pos"): val for key, val in action.items() if key.endswith(".pos")}
|
||||
|
||||
# Apply joint limit clipping to arm
|
||||
@@ -343,7 +349,6 @@ class OpenArmFollower(Robot):
|
||||
@check_if_not_connected
|
||||
def disconnect(self):
|
||||
"""Disconnect from robot."""
|
||||
|
||||
# Disconnect CAN bus
|
||||
self.bus.disconnect(self.config.disable_torque_on_disconnect)
|
||||
|
||||
|
||||
@@ -23,6 +23,58 @@ from ..config import RobotConfig
|
||||
@RobotConfig.register_subclass("reachy2")
|
||||
@dataclass
|
||||
class Reachy2RobotConfig(RobotConfig):
|
||||
"""Configuration for the Reachy 2 humanoid.
|
||||
|
||||
Reachy 2 is driven over the network rather than a serial bus, so `port` is a TCP port on the robot's
|
||||
gRPC service rather than a device path. Calibration is handled by the robot itself and there is no
|
||||
LeRobot calibration file.
|
||||
|
||||
Which joints appear in observations and actions is selected by the `with_*` flags: turning a part off
|
||||
removes its joints entirely. At least one part must stay enabled.
|
||||
|
||||
Args:
|
||||
max_relative_target (`float`, *optional*):
|
||||
Caps how far a single action may move a joint from its present position, as a safety limit.
|
||||
`None` disables clipping.
|
||||
ip_address (`str`, *optional*, defaults to `"localhost"`):
|
||||
Address of the Reachy 2 robot.
|
||||
port (`int`, *optional*, defaults to 50065):
|
||||
TCP port of the robot's service. Not a serial port.
|
||||
disable_torque_on_disconnect (`bool`, *optional*, defaults to `False`):
|
||||
Whether to call `turn_off_smoothly()` before disconnecting.
|
||||
use_external_commands (`bool`, *optional*, defaults to `False`):
|
||||
Set `True` when another system drives the robot, such as the official
|
||||
[teleoperation app](https://github.com/pollen-robotics/Reachy2Teleoperation). In that mode
|
||||
[`~robots.Robot.send_action`] does not send anything to the robot.
|
||||
with_mobile_base (`bool`, *optional*, defaults to `True`):
|
||||
Whether to include the mobile base's joints.
|
||||
with_l_arm (`bool`, *optional*, defaults to `True`):
|
||||
Whether to include the left arm's joints.
|
||||
with_r_arm (`bool`, *optional*, defaults to `True`):
|
||||
Whether to include the right arm's joints.
|
||||
with_neck (`bool`, *optional*, defaults to `True`):
|
||||
Whether to include the neck's joints.
|
||||
with_antennas (`bool`, *optional*, defaults to `True`):
|
||||
Whether to include the antennas' joints.
|
||||
with_left_teleop_camera (`bool`, *optional*, defaults to `False`):
|
||||
Whether to add the left teleoperation camera to observations.
|
||||
with_right_teleop_camera (`bool`, *optional*, defaults to `False`):
|
||||
Whether to add the right teleoperation camera to observations.
|
||||
with_torso_camera (`bool`, *optional*, defaults to `False`):
|
||||
Whether to add the torso RGB camera to observations.
|
||||
camera_width (`int`, *optional*, defaults to 640):
|
||||
Frame width for the built-in cameras. Their frame rate is fixed at 30 and is not configurable.
|
||||
camera_height (`int`, *optional*, defaults to 480):
|
||||
Frame height for the built-in cameras.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
Additional cameras beyond the three built-in ones. The `with_*_camera` flags populate this
|
||||
field, so anything set here is merged with them.
|
||||
id (`str`, *optional*):
|
||||
Identifier for this particular robot.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Unused: Reachy 2 manages its own calibration.
|
||||
"""
|
||||
|
||||
# `max_relative_target` limits the magnitude of the relative positional target vector for safety purposes.
|
||||
# Set this to a positive scalar to have the same value for all motors.
|
||||
max_relative_target: float | None = None
|
||||
@@ -65,6 +117,11 @@ class Reachy2RobotConfig(RobotConfig):
|
||||
cameras: dict[str, CameraConfig] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""Add the built-in cameras selected by the `with_*_camera` flags and validate the part selection.
|
||||
|
||||
Raises:
|
||||
ValueError: If every robot part is disabled, which would leave no joints to control.
|
||||
"""
|
||||
# Add cameras with same ip_address as the robot
|
||||
if self.with_left_teleop_camera:
|
||||
self.cameras["teleop_left"] = Reachy2CameraConfig(
|
||||
|
||||
@@ -73,14 +73,18 @@ REACHY2_VEL = {
|
||||
|
||||
|
||||
class Reachy2Robot(Robot):
|
||||
"""
|
||||
[Reachy 2](https://www.pollen-robotics.com/reachy/), by Pollen Robotics.
|
||||
"""
|
||||
"""[Reachy 2](https://www.pollen-robotics.com/reachy/), the humanoid by Pollen Robotics."""
|
||||
|
||||
config_class = Reachy2RobotConfig
|
||||
name = "reachy2"
|
||||
|
||||
def __init__(self, config: Reachy2RobotConfig):
|
||||
"""Build the robot from its configuration.
|
||||
|
||||
Args:
|
||||
config (`Reachy2RobotConfig`):
|
||||
The robot's configuration. Its `port` and `cameras` determine what is connected.
|
||||
"""
|
||||
require_package("reachy2_sdk", extra="reachy2")
|
||||
super().__init__(config)
|
||||
|
||||
@@ -97,18 +101,41 @@ class Reachy2Robot(Robot):
|
||||
|
||||
@property
|
||||
def observation_features(self) -> dict[str, Any]:
|
||||
"""The values this robot reports, and their types or shapes.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys as returned by [`~robots.Robot.get_observation`], mapped to a scalar type for
|
||||
proprioceptive values or to a `(height, width, channels)` shape for images.
|
||||
"""
|
||||
return {**self.motors_features, **self.camera_features}
|
||||
|
||||
@property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
"""The values this robot accepts, and their types.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys accepted by [`~robots.Robot.send_action`], mapped to their type.
|
||||
"""
|
||||
return self.motors_features
|
||||
|
||||
@property
|
||||
def camera_features(self) -> dict[str, tuple[int | None, int | None, int]]:
|
||||
"""The shape of each configured camera's frames.
|
||||
|
||||
Returns:
|
||||
`dict[str, tuple[int | None, int | None, int]]`: Camera name mapped to
|
||||
`(height, width, channels)`.
|
||||
"""
|
||||
return {cam: (self.cameras[cam].height, self.cameras[cam].width, 3) for cam in self.cameras}
|
||||
|
||||
@property
|
||||
def motors_features(self) -> dict[str, type]:
|
||||
"""The joints this robot exposes, given which parts are enabled in the config.
|
||||
|
||||
Returns:
|
||||
`dict[str, type]`: Joint name mapped to `float`, including the mobile base's velocity
|
||||
components when `with_mobile_base` is set.
|
||||
"""
|
||||
if self.config.with_mobile_base:
|
||||
return {
|
||||
**dict.fromkeys(
|
||||
@@ -125,9 +152,23 @@ class Reachy2Robot(Robot):
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Whether every device this robot uses is connected.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` only when the robot and all its cameras are connected.
|
||||
"""
|
||||
return self.reachy.is_connected() if self.reachy is not None else False
|
||||
|
||||
def connect(self, calibrate: bool = False) -> None:
|
||||
"""Connect to the robot and its cameras, then apply the configured settings.
|
||||
|
||||
Args:
|
||||
calibrate (`bool`, *optional*, defaults to `False`):
|
||||
Accepted for interface compatibility and ignored: Reachy 2 manages its own calibration.
|
||||
|
||||
Raises:
|
||||
DeviceAlreadyConnectedError: If the robot is already connected.
|
||||
"""
|
||||
self.reachy = ReachySDK(self.config.ip_address)
|
||||
if not self.is_connected:
|
||||
raise ConnectionError()
|
||||
@@ -138,15 +179,25 @@ class Reachy2Robot(Robot):
|
||||
self.configure()
|
||||
|
||||
def configure(self) -> None:
|
||||
"""Apply the operating mode, gains and limits from the configuration to the robot."""
|
||||
if self.reachy is not None:
|
||||
self.reachy.turn_on()
|
||||
self.reachy.reset_default_limits()
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
"""Whether the robot is calibrated.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` when no calibration is needed before use.
|
||||
"""
|
||||
return True
|
||||
|
||||
def calibrate(self) -> None:
|
||||
"""Calibrate the robot and store the result.
|
||||
|
||||
Interactive: prompts on stdin and asks you to move the robot through the required positions.
|
||||
"""
|
||||
pass
|
||||
|
||||
def _generate_joints_dict(self) -> dict[str, str]:
|
||||
@@ -172,6 +223,14 @@ class Reachy2Robot(Robot):
|
||||
return {}
|
||||
|
||||
def get_observation(self) -> RobotObservation:
|
||||
"""Read the robot's current state and a frame from each camera.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: Keys matching [`~robots.Robot.observation_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
obs_dict: RobotObservation = {}
|
||||
|
||||
# Read Reachy 2 state
|
||||
@@ -186,6 +245,18 @@ class Reachy2Robot(Robot):
|
||||
return obs_dict
|
||||
|
||||
def send_action(self, action: RobotAction) -> RobotAction:
|
||||
"""Command the robot to move towards a target configuration.
|
||||
|
||||
Args:
|
||||
action (`dict[str, Any]`):
|
||||
Target values, keyed as in [`~robots.Robot.action_features`].
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The action actually sent, which may be clipped by `max_relative_target`.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
if self.reachy is not None:
|
||||
if not self.is_connected:
|
||||
raise ConnectionError()
|
||||
@@ -228,6 +299,11 @@ class Reachy2Robot(Robot):
|
||||
return action
|
||||
|
||||
def disconnect(self) -> None:
|
||||
"""Disconnect from the robot and its cameras.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
if self.reachy is not None:
|
||||
for cam in self.cameras.values():
|
||||
cam.disconnect()
|
||||
|
||||
@@ -23,10 +23,12 @@ from ..config import RobotConfig
|
||||
|
||||
@dataclass
|
||||
class RebotB601FollowerConfig:
|
||||
"""Base configuration class for the Seeed Studio reBot B601-DM follower arm.
|
||||
"""Field definitions for the Seeed Studio reBot B601-DM follower arm.
|
||||
|
||||
The B601-DM is a 6-DOF arm plus gripper driven by Damiao CAN motors. Motor
|
||||
communication goes through the ``motorbridge`` package.
|
||||
This class only carries the fields. The registered configuration users instantiate is
|
||||
[`RebotB601FollowerRobotConfig`], which documents them all in one place — doc-builder renders only a
|
||||
class's own docstring, never its bases'. It is also used directly as the per-arm config of
|
||||
[`~robots.bi_rebot_b601_follower.BiRebotB601FollowerConfig`].
|
||||
"""
|
||||
|
||||
# Communication port. For ``can_adapter="damiao"`` this is the Damiao serial
|
||||
@@ -104,6 +106,62 @@ class RebotB601FollowerConfig:
|
||||
@RobotConfig.register_subclass("rebot_b601_follower")
|
||||
@dataclass
|
||||
class RebotB601FollowerRobotConfig(RobotConfig, RebotB601FollowerConfig):
|
||||
"""Registered configuration for the reBot B601-DM follower robot."""
|
||||
"""Configuration for the Seeed Studio reBot B601-DM follower arm.
|
||||
|
||||
The B601-DM is a 6-DOF arm plus gripper on Damiao CAN motors, driven through the `motorbridge`
|
||||
package. What `port` means depends on `can_adapter`. Calibration follows the usual LeRobot flow and is
|
||||
stored per `id`.
|
||||
|
||||
The arm and the gripper are controlled separately: `control_mode` governs the six arm joints and
|
||||
`gripper_control_mode` the gripper, and each mode uses a different subset of the gain fields.
|
||||
|
||||
Per-joint lists hold 7 values in motor order: `shoulder_pan`, `shoulder_lift`, `elbow_flex`,
|
||||
`wrist_flex`, `wrist_yaw`, `wrist_roll`, `gripper`.
|
||||
|
||||
Args:
|
||||
port (`str`):
|
||||
Where the arm is reached. For `can_adapter="damiao"` this is the serial bridge device, e.g.
|
||||
`/dev/ttyACM0`; for `can_adapter="socketcan"` it is the CAN channel name, e.g. `can0`.
|
||||
can_adapter (`str`, *optional*, defaults to `"damiao"`):
|
||||
`"damiao"` for the dedicated Damiao serial bridge, or `"socketcan"` for SocketCAN adapters
|
||||
such as PCAN, slcan and embedded controllers.
|
||||
dm_serial_baud (`int`, *optional*, defaults to 921600):
|
||||
Baud rate of the Damiao serial bridge. Only used when `can_adapter="damiao"`.
|
||||
disable_torque_on_disconnect (`bool`, *optional*, defaults to `True`):
|
||||
Whether to release the motors on disconnect. Leave `True` unless the arm is holding a load it
|
||||
must not drop.
|
||||
max_relative_target (`float | dict[str, float]`, *optional*):
|
||||
Caps how far a single action may move the arm from its present position, in degrees. A scalar
|
||||
applies to every motor; a dict maps motor name to a per-motor cap. `None` disables clipping.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
Cameras to read alongside the arm's joint positions.
|
||||
motor_can_ids (`dict[str, tuple[int, int]]`, *optional*):
|
||||
Maps motor name to its `(send_can_id, recv_can_id)` pair. Change it only if you have
|
||||
re-addressed the motors.
|
||||
pos_vel_velocity (`float | list[float]`, *optional*):
|
||||
Maximum speed in deg/s per joint, used by the arm in `pos_vel` mode and by the gripper in
|
||||
`force_pos` mode.
|
||||
control_mode (`str`, *optional*, defaults to `"mit"`):
|
||||
How the six arm joints are driven: `"mit"` or `"pos_vel"`.
|
||||
mit_kp (`float | list[float]`, *optional*):
|
||||
MIT-mode proportional gains per arm joint. Unused when `control_mode="pos_vel"`.
|
||||
mit_kd (`float | list[float]`, *optional*):
|
||||
MIT-mode derivative gains per arm joint. Unused when `control_mode="pos_vel"`.
|
||||
gripper_control_mode (`str`, *optional*, defaults to `"force_pos"`):
|
||||
How the gripper is driven: `"force_pos"` or `"mit"`.
|
||||
gripper_torque_ratio (`float`, *optional*, defaults to 0.07):
|
||||
Maximum grip force as a fraction in `[0, 1]`. Only used when
|
||||
`gripper_control_mode="force_pos"`.
|
||||
gripper_mit_kp (`float`, *optional*, defaults to 8.0):
|
||||
Gripper MIT-mode proportional gain. Only used when `gripper_control_mode="mit"`.
|
||||
gripper_mit_kd (`float`, *optional*, defaults to 0.3):
|
||||
Gripper MIT-mode derivative gain. Only used when `gripper_control_mode="mit"`.
|
||||
joint_limits (`dict[str, tuple[float, float]]`, *optional*):
|
||||
Soft `(min, max)` limits in degrees per joint, clipped against on every action.
|
||||
id (`str`, *optional*):
|
||||
Identifier for this particular arm; also names its calibration file.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
@@ -66,6 +66,12 @@ class RebotB601Follower(Robot):
|
||||
name = "rebot_b601_follower"
|
||||
|
||||
def __init__(self, config: RebotB601FollowerRobotConfig):
|
||||
"""Build the robot from its configuration.
|
||||
|
||||
Args:
|
||||
config (`RebotB601FollowerRobotConfig`):
|
||||
The robot's configuration. Its `port` and `cameras` determine what is connected.
|
||||
"""
|
||||
require_package("motorbridge", extra="rebot")
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
@@ -91,18 +97,43 @@ class RebotB601Follower(Robot):
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
"""The values this robot reports, and their types or shapes.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys as returned by [`~robots.Robot.get_observation`], mapped to a scalar type for
|
||||
proprioceptive values or to a `(height, width, channels)` shape for images.
|
||||
"""
|
||||
return {**self._motors_ft, **self._cameras_ft}
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
"""The values this robot accepts, and their types.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys accepted by [`~robots.Robot.send_action`], mapped to their type.
|
||||
"""
|
||||
return self._motors_ft
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Whether every device this robot uses is connected.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` only when the robot and all its cameras are connected.
|
||||
"""
|
||||
return self.bus is not None and all(cam.is_connected for cam in self.cameras.values())
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""Connect to the robot and its cameras, then apply the configured settings.
|
||||
|
||||
Args:
|
||||
calibrate (`bool`, *optional*, defaults to `True`):
|
||||
Whether to run calibration if the robot is not already calibrated.
|
||||
|
||||
Raises:
|
||||
DeviceAlreadyConnectedError: If the robot is already connected.
|
||||
"""
|
||||
logger.info(f"Connecting {self} on {self.config.port} (adapter={self.config.can_adapter})...")
|
||||
if self.config.can_adapter == "damiao":
|
||||
self.bus = MotorBridgeController.from_dm_serial(
|
||||
@@ -133,9 +164,18 @@ class RebotB601Follower(Robot):
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
"""Whether the robot is calibrated.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` when no calibration is needed before use.
|
||||
"""
|
||||
return bool(self.calibration)
|
||||
|
||||
def calibrate(self) -> None:
|
||||
"""Calibrate the robot and store the result.
|
||||
|
||||
Interactive: prompts on stdin and asks you to move the robot through the required positions.
|
||||
"""
|
||||
if self.calibration:
|
||||
user_input = input(
|
||||
f"Press ENTER to use provided calibration file associated with the id {self.id}, "
|
||||
@@ -174,6 +214,7 @@ class RebotB601Follower(Robot):
|
||||
print(f"Calibration saved to {self.calibration_fpath}")
|
||||
|
||||
def configure(self) -> None:
|
||||
"""Apply the operating mode, gains and limits from the configuration to the robot."""
|
||||
if self.config.control_mode not in ("pos_vel", "mit"):
|
||||
raise ValueError(
|
||||
f"Unsupported control_mode '{self.config.control_mode}'. Use 'pos_vel' or 'mit'."
|
||||
@@ -226,6 +267,14 @@ class RebotB601Follower(Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
"""Read the robot's current state and a frame from each camera.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: Keys matching [`~robots.Robot.observation_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
start = time.perf_counter()
|
||||
obs_dict = {f"{motor}.pos": pos for motor, pos in self._present_pos().items()}
|
||||
dt_ms = (time.perf_counter() - start) * 1e3
|
||||
@@ -311,6 +360,11 @@ class RebotB601Follower(Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self) -> None:
|
||||
"""Disconnect from the robot and its cameras.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
for motor in self.motors.values():
|
||||
if self.config.disable_torque_on_disconnect:
|
||||
motor.disable()
|
||||
|
||||
+95
-63
@@ -28,15 +28,22 @@ from .config import RobotConfig
|
||||
# TODO(aliberts): action/obs typing such as Generic[ObsType, ActType] similar to gym.Env ?
|
||||
# https://github.com/Farama-Foundation/Gymnasium/blob/3287c869f9a48d99454306b0d4b4ec537f0f35e3/gymnasium/core.py#L23
|
||||
class Robot(abc.ABC):
|
||||
"""
|
||||
The base abstract class for all LeRobot-compatible robots.
|
||||
"""The base abstract class for all LeRobot-compatible robots.
|
||||
|
||||
This class provides a standardized interface for interacting with physical robots.
|
||||
Subclasses must implement all abstract methods and properties to be usable.
|
||||
This class provides a standardized interface for interacting with physical robots. Subclasses must
|
||||
implement all abstract methods and properties to be usable.
|
||||
|
||||
Attributes:
|
||||
config_class (RobotConfig): The expected configuration class for this robot.
|
||||
name (str): The unique robot name used to identify this robot type.
|
||||
Used as a context manager, a robot connects on entry and disconnects on exit even if the body raises:
|
||||
|
||||
```python
|
||||
>>> with SO101Follower(config) as robot: # doctest: +SKIP
|
||||
... obs = robot.get_observation()
|
||||
... robot.send_action(action)
|
||||
```
|
||||
|
||||
**Attributes**:
|
||||
- **config_class** (`type[RobotConfig]`) -- The expected configuration class for this robot.
|
||||
- **name** (`str`) -- The unique robot name used to identify this robot type.
|
||||
"""
|
||||
|
||||
# Set these in ALL subclasses
|
||||
@@ -44,6 +51,13 @@ class Robot(abc.ABC):
|
||||
name: str
|
||||
|
||||
def __init__(self, config: RobotConfig):
|
||||
"""Set up identity and calibration paths, loading an existing calibration file if there is one.
|
||||
|
||||
Args:
|
||||
config (`RobotConfig`):
|
||||
The robot's configuration. Its `id` and `calibration_dir` decide where calibration is
|
||||
read from and written to.
|
||||
"""
|
||||
self.robot_type = self.name
|
||||
self.id = config.id
|
||||
self.calibration_dir = (
|
||||
@@ -56,28 +70,24 @@ class Robot(abc.ABC):
|
||||
self._load_calibration()
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Return this robot's id and class name, e.g. `"my_arm SO101Follower"`.
|
||||
|
||||
Returns:
|
||||
`str`: A short identifier used in log messages.
|
||||
"""
|
||||
return f"{self.id} {self.__class__.__name__}"
|
||||
|
||||
def __enter__(self):
|
||||
"""
|
||||
Context manager entry.
|
||||
Automatically connects to the camera.
|
||||
"""
|
||||
"""Context manager entry. Automatically connects to the robot."""
|
||||
self.connect()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
||||
"""
|
||||
Context manager exit.
|
||||
Automatically disconnects, ensuring resources are released even on error.
|
||||
"""
|
||||
"""Context manager exit. Disconnects, ensuring resources are released even on error."""
|
||||
self.disconnect()
|
||||
|
||||
def __del__(self) -> None:
|
||||
"""
|
||||
Destructor safety net.
|
||||
Attempts to disconnect if the object is garbage collected without cleanup.
|
||||
"""
|
||||
"""Destructor safety net. Disconnects if the object is garbage collected without cleanup."""
|
||||
try:
|
||||
if self.is_connected:
|
||||
self.disconnect()
|
||||
@@ -88,83 +98,102 @@ class Robot(abc.ABC):
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def observation_features(self) -> dict:
|
||||
"""
|
||||
A dictionary describing the structure and types of the observations produced by the robot.
|
||||
Its structure (keys) should match the structure of what is returned by :pymeth:`get_observation`.
|
||||
Values for the dict should either be:
|
||||
- The type of the value if it's a simple value, e.g. `float` for single proprioceptive value (a joint's position/velocity)
|
||||
- A tuple representing the shape if it's an array-type value, e.g. `(height, width, channel)` for images
|
||||
"""A dictionary describing the structure and types of the observations produced by the robot.
|
||||
|
||||
Note: this property should be able to be called regardless of whether the robot is connected or not.
|
||||
Its keys should match the structure of what is returned by [`~robots.Robot.get_observation`]. Values
|
||||
should either be:
|
||||
|
||||
- the type of the value if it's a simple value, e.g. `float` for a single proprioceptive value
|
||||
(a joint's position or velocity)
|
||||
- a tuple representing the shape if it's an array-type value, e.g. `(height, width, channel)` for
|
||||
images
|
||||
|
||||
> [!NOTE]
|
||||
> This property must be callable regardless of whether the robot is connected.
|
||||
|
||||
Returns:
|
||||
`dict`: Observation names mapped to their type or shape.
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def action_features(self) -> dict:
|
||||
"""
|
||||
A dictionary describing the structure and types of the actions expected by the robot. Its structure
|
||||
(keys) should match the structure of what is passed to :pymeth:`send_action`. Values for the dict
|
||||
should be the type of the value if it's a simple value, e.g. `float` for single proprioceptive value
|
||||
(a joint's goal position/velocity)
|
||||
"""A dictionary describing the structure and types of the actions expected by the robot.
|
||||
|
||||
Note: this property should be able to be called regardless of whether the robot is connected or not.
|
||||
Its keys should match the structure of what is passed to [`~robots.Robot.send_action`]. Values should
|
||||
be the type of the value if it's a simple value, e.g. `float` for a single proprioceptive value
|
||||
(a joint's goal position or velocity).
|
||||
|
||||
> [!NOTE]
|
||||
> This property must be callable regardless of whether the robot is connected.
|
||||
|
||||
Returns:
|
||||
`dict`: Action names mapped to their type or shape.
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def is_connected(self) -> bool:
|
||||
"""
|
||||
Whether the robot is currently connected or not. If `False`, calling :pymeth:`get_observation` or
|
||||
:pymeth:`send_action` should raise an error.
|
||||
"""Whether the robot is currently connected.
|
||||
|
||||
If `False`, calling [`~robots.Robot.get_observation`] or [`~robots.Robot.send_action`] should raise
|
||||
an error.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` if communication with the robot is established.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""
|
||||
Establish communication with the robot.
|
||||
"""Establish communication with the robot.
|
||||
|
||||
Args:
|
||||
calibrate (bool): If True, automatically calibrate the robot after connecting if it's not
|
||||
calibrated or needs calibration (this is hardware-dependant).
|
||||
calibrate (`bool`, *optional*, defaults to `True`):
|
||||
Whether to automatically calibrate the robot after connecting, if it is not calibrated or
|
||||
needs recalibration. Whether calibration is needed is hardware-dependent.
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def is_calibrated(self) -> bool:
|
||||
"""Whether the robot is currently calibrated or not. Should be always `True` if not applicable"""
|
||||
"""Whether the robot is currently calibrated.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` if the robot is calibrated. Always `True` for robots where calibration does not
|
||||
apply.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def calibrate(self) -> None:
|
||||
"""
|
||||
Calibrate the robot if applicable. If not, this should be a no-op.
|
||||
"""Calibrate the robot if applicable. If not, this should be a no-op.
|
||||
|
||||
This method should collect any necessary data (e.g., motor offsets) and update the
|
||||
:pyattr:`calibration` dictionary accordingly.
|
||||
This method should collect any necessary data (e.g. motor offsets) and update the `calibration`
|
||||
dictionary accordingly.
|
||||
"""
|
||||
pass
|
||||
|
||||
def _load_calibration(self, fpath: Path | None = None) -> None:
|
||||
"""
|
||||
Helper to load calibration data from the specified file.
|
||||
"""Helper to load calibration data from the specified file.
|
||||
|
||||
Args:
|
||||
fpath (Path | None): Optional path to the calibration file. Defaults to `self.calibration_fpath`.
|
||||
fpath (`Path`, *optional*):
|
||||
Path to the calibration file. Defaults to `self.calibration_fpath`.
|
||||
"""
|
||||
fpath = self.calibration_fpath if fpath is None else fpath
|
||||
with open(fpath) as f, draccus.config_type("json"):
|
||||
self.calibration = draccus.load(dict[str, MotorCalibration], f)
|
||||
|
||||
def _save_calibration(self, fpath: Path | None = None) -> None:
|
||||
"""
|
||||
Helper to save calibration data to the specified file.
|
||||
"""Helper to save calibration data to the specified file.
|
||||
|
||||
Args:
|
||||
fpath (Path | None): Optional path to save the calibration file. Defaults to `self.calibration_fpath`.
|
||||
fpath (`Path`, *optional*):
|
||||
Path to save the calibration file to. Defaults to `self.calibration_fpath`.
|
||||
"""
|
||||
fpath = self.calibration_fpath if fpath is None else fpath
|
||||
with open(fpath, "w") as f, draccus.config_type("json"):
|
||||
@@ -172,36 +201,39 @@ class Robot(abc.ABC):
|
||||
|
||||
@abc.abstractmethod
|
||||
def configure(self) -> None:
|
||||
"""
|
||||
Apply any one-time or runtime configuration to the robot.
|
||||
"""Apply any one-time or runtime configuration to the robot.
|
||||
|
||||
This may include setting motor parameters, control modes, or initial state.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def get_observation(self) -> RobotObservation:
|
||||
"""
|
||||
Retrieve the current observation from the robot.
|
||||
"""Retrieve the current observation from the robot.
|
||||
|
||||
Returns:
|
||||
RobotObservation: A flat dictionary representing the robot's current sensory state. Its structure
|
||||
should match :pymeth:`observation_features`.
|
||||
"""
|
||||
`dict[str, Any]`: A flat dictionary representing the robot's current sensory state. Its structure
|
||||
should match [`~robots.Robot.observation_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If [`~robots.Robot.connect`] has not been called.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def send_action(self, action: RobotAction) -> RobotAction:
|
||||
"""
|
||||
Send an action command to the robot.
|
||||
"""Send an action command to the robot.
|
||||
|
||||
Args:
|
||||
action (RobotAction): Dictionary representing the desired action. Its structure should match
|
||||
:pymeth:`action_features`.
|
||||
action (`dict[str, Any]`):
|
||||
The desired action. Its structure should match [`~robots.Robot.action_features`].
|
||||
|
||||
Returns:
|
||||
RobotAction: The action actually sent to the motors potentially clipped or modified, e.g. by
|
||||
safety limits on velocity.
|
||||
`dict[str, Any]`: The action actually sent to the motors, potentially clipped or modified, e.g.
|
||||
by safety limits on velocity. Prefer this over the requested action when logging or recording.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If [`~robots.Robot.connect`] has not been called.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@@ -23,7 +23,12 @@ from ..config import RobotConfig
|
||||
|
||||
@dataclass
|
||||
class SOFollowerConfig:
|
||||
"""Base configuration class for SO Follower robots."""
|
||||
"""Field definitions shared by the SO-family follower arms.
|
||||
|
||||
This class only carries the fields. The registered configuration users instantiate is
|
||||
[`SOFollowerRobotConfig`], which combines these with [`~robots.RobotConfig`] and documents them all in
|
||||
one place — doc-builder renders only a class's own docstring, never its bases'.
|
||||
"""
|
||||
|
||||
# Port to connect to the arm
|
||||
port: str
|
||||
@@ -57,6 +62,51 @@ class SOFollowerConfig:
|
||||
@RobotConfig.register_subclass("so100_follower")
|
||||
@dataclass
|
||||
class SOFollowerRobotConfig(RobotConfig, SOFollowerConfig):
|
||||
"""Configuration for the SO-100 and SO-101 follower arms.
|
||||
|
||||
Both arms share this class; `SO100FollowerConfig` and `SO101FollowerConfig` are aliases for it. They
|
||||
differ in their calibration and gearing, not in their control code.
|
||||
|
||||
Args:
|
||||
port (`str`):
|
||||
Serial port the arm is connected to, e.g. `/dev/ttyACM0` on Linux or `COM3` on Windows. Run
|
||||
`lerobot-find-port` to identify it.
|
||||
disable_torque_on_disconnect (`bool`, *optional*, defaults to `True`):
|
||||
Whether to release the motors on disconnect. Leave `True` unless the arm is holding a load it
|
||||
must not drop.
|
||||
max_relative_target (`float | dict[str, float]`, *optional*):
|
||||
Caps how far a single action may move the arm from its present position, as a safety limit. A
|
||||
scalar applies to every motor; a dict maps motor name to a per-motor cap. `None` disables
|
||||
clipping. Enabling this costs an extra read of the present position on every step.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
Cameras to read alongside the arm's joint positions, keyed by the name they appear under in
|
||||
observations. Each must specify `width`, `height` and `fps`.
|
||||
use_degrees (`bool`, *optional*, defaults to `True`):
|
||||
Whether to report and accept joint positions in degrees. Keep `True` for compatibility with
|
||||
existing policies and datasets.
|
||||
position_p_coefficient (`int`, *optional*, defaults to 16):
|
||||
Proportional gain written to the Feetech STS3215 motors at connect time.
|
||||
position_i_coefficient (`int`, *optional*, defaults to 0):
|
||||
Integral gain written to the motors at connect time.
|
||||
position_d_coefficient (`int`, *optional*, defaults to 32):
|
||||
Derivative gain written to the motors at connect time.
|
||||
num_read_retries (`int`, *optional*, defaults to 2):
|
||||
Extra attempts when a `sync_read` fails. Feetech buses occasionally return a corrupted status
|
||||
packet, especially when several joints move at once, which would otherwise abort the control
|
||||
loop. Retries are immediate and only happen on failure, so steady-state read cost is unchanged.
|
||||
id (`str`, *optional*):
|
||||
Identifier for this particular arm; also names its calibration file.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Where to read and write the calibration file. Defaults to the LeRobot calibration home.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> from lerobot.robots.so_follower import SO101Follower, SO101FollowerConfig
|
||||
>>> config = SO101FollowerConfig(port="/dev/ttyACM0", max_relative_target=5.0) # doctest: +SKIP
|
||||
>>> robot = SO101Follower(config) # doctest: +SKIP
|
||||
```
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
|
||||
@@ -40,8 +40,7 @@ logger = logging.getLogger(__name__)
|
||||
@ProcessorStepRegistry.register("ee_reference_and_delta")
|
||||
@dataclass
|
||||
class EEReferenceAndDelta(RobotActionProcessorStep):
|
||||
"""
|
||||
Computes a target end-effector pose from a relative delta command.
|
||||
"""Computes a target end-effector pose from a relative delta command.
|
||||
|
||||
This step takes a desired change in position and orientation (`target_*`) and applies it to a
|
||||
reference end-effector pose to calculate an absolute target pose. The reference pose is derived
|
||||
@@ -53,15 +52,16 @@ class EEReferenceAndDelta(RobotActionProcessorStep):
|
||||
2. `use_latched_reference=False`: The reference pose is updated to the robot's current pose at
|
||||
every step.
|
||||
|
||||
Attributes:
|
||||
kinematics: The robot's kinematic model for forward kinematics.
|
||||
end_effector_step_sizes: A dictionary scaling the input delta commands.
|
||||
motor_names: A list of motor names required for forward kinematics.
|
||||
use_latched_reference: If True, latch the reference pose on enable; otherwise, always use the
|
||||
current pose as the reference.
|
||||
reference_ee_pose: Internal state storing the latched reference pose.
|
||||
_prev_enabled: Internal state to detect the rising edge of the enable signal.
|
||||
_command_when_disabled: Internal state to hold the last command while disabled.
|
||||
**Attributes**:
|
||||
- **kinematics** (`RobotKinematics`) -- The robot's kinematic model for forward kinematics.
|
||||
- **end_effector_step_sizes** (`dict`) -- A dictionary scaling the input delta commands.
|
||||
- **motor_names** (`list[str]`) -- A list of motor names required for forward kinematics.
|
||||
- **use_latched_reference** (`bool`) -- If True, latch the reference pose on enable; otherwise, always
|
||||
use the current pose as the reference.
|
||||
- **reference_ee_pose** (`np.ndarray | None`) -- Internal state storing the latched reference pose.
|
||||
- **_prev_enabled** (`bool`) -- Internal state to detect the rising edge of the enable signal.
|
||||
- **_command_when_disabled** (`np.ndarray | None`) -- Internal state to hold the last command while
|
||||
disabled.
|
||||
"""
|
||||
|
||||
kinematics: RobotKinematics
|
||||
@@ -77,6 +77,15 @@ class EEReferenceAndDelta(RobotActionProcessorStep):
|
||||
_command_when_disabled: np.ndarray | None = field(default=None, init=False, repr=False)
|
||||
|
||||
def action(self, action: RobotAction) -> RobotAction:
|
||||
"""Transform the action for this step.
|
||||
|
||||
Args:
|
||||
action (`dict[str, Any]`):
|
||||
The incoming robot action.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The transformed action.
|
||||
"""
|
||||
raw_observation = self.transition.get(TransitionKey.OBSERVATION)
|
||||
|
||||
if raw_observation is None:
|
||||
@@ -167,6 +176,16 @@ class EEReferenceAndDelta(RobotActionProcessorStep):
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
"""Update the feature contract to match what this step does to the data.
|
||||
|
||||
Args:
|
||||
features (`dict[PipelineFeatureType, dict[str, PolicyFeature]]`):
|
||||
The pipeline's feature contract so far.
|
||||
|
||||
Returns:
|
||||
`dict[PipelineFeatureType, dict[str, PolicyFeature]]`: The contract with this step's key changes
|
||||
applied.
|
||||
"""
|
||||
for feat in [
|
||||
"enabled",
|
||||
"target_x",
|
||||
@@ -190,21 +209,19 @@ class EEReferenceAndDelta(RobotActionProcessorStep):
|
||||
@ProcessorStepRegistry.register("ee_bounds_and_safety")
|
||||
@dataclass
|
||||
class EEBoundsAndSafety(RobotActionProcessorStep):
|
||||
"""
|
||||
Clips the end-effector pose to predefined bounds and checks for unsafe jumps.
|
||||
"""Clips the end-effector pose to predefined bounds and checks for unsafe jumps.
|
||||
|
||||
This step ensures that the target end-effector pose remains within a safe operational workspace.
|
||||
It also moderates the command to prevent large, sudden movements between consecutive steps.
|
||||
|
||||
Attributes:
|
||||
end_effector_bounds: A dictionary with "min" and "max" keys for position clipping.
|
||||
max_ee_step_m: The maximum allowed change in position (in meters) between steps.
|
||||
raise_on_jump: When ``True`` (default) an over-limit per-frame step raises
|
||||
``ValueError`` (aborting the control loop). When ``False`` the step is
|
||||
rate-limited to ``max_ee_step_m`` and a warning is logged instead — the
|
||||
safer choice for live teleoperation, where a transient tracking glitch
|
||||
should not crash the loop and leave the robot uncontrolled.
|
||||
_last_pos: Internal state storing the last commanded position.
|
||||
**Attributes**:
|
||||
- **end_effector_bounds** (`dict`) -- A dictionary with "min" and "max" keys for position clipping.
|
||||
- **max_ee_step_m** (`float`) -- The maximum allowed change in position (in meters) between steps.
|
||||
- **raise_on_jump** (`bool`) -- When ``True`` (default) an over-limit per-frame step raises
|
||||
``ValueError`` (aborting the control loop). When ``False`` the step is rate-limited to
|
||||
``max_ee_step_m`` and a warning is logged instead — the safer choice for live teleoperation, where a
|
||||
transient tracking glitch should not crash the loop and leave the robot uncontrolled.
|
||||
- **_last_pos** (`np.ndarray | None`) -- Internal state storing the last commanded position.
|
||||
"""
|
||||
|
||||
end_effector_bounds: dict
|
||||
@@ -213,6 +230,15 @@ class EEBoundsAndSafety(RobotActionProcessorStep):
|
||||
_last_pos: np.ndarray | None = field(default=None, init=False, repr=False)
|
||||
|
||||
def action(self, action: RobotAction) -> RobotAction:
|
||||
"""Transform the action for this step.
|
||||
|
||||
Args:
|
||||
action (`dict[str, Any]`):
|
||||
The incoming robot action.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The transformed action.
|
||||
"""
|
||||
x = action["ee.x"]
|
||||
y = action["ee.y"]
|
||||
z = action["ee.z"]
|
||||
@@ -268,29 +294,39 @@ class EEBoundsAndSafety(RobotActionProcessorStep):
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
"""Update the feature contract to match what this step does to the data.
|
||||
|
||||
Args:
|
||||
features (`dict[PipelineFeatureType, dict[str, PolicyFeature]]`):
|
||||
The pipeline's feature contract so far.
|
||||
|
||||
Returns:
|
||||
`dict[PipelineFeatureType, dict[str, PolicyFeature]]`: The contract with this step's key changes
|
||||
applied.
|
||||
"""
|
||||
return features
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register("inverse_kinematics_ee_to_joints")
|
||||
@dataclass
|
||||
class InverseKinematicsEEToJoints(RobotActionProcessorStep):
|
||||
"""
|
||||
Computes desired joint positions from a target end-effector pose using inverse kinematics (IK).
|
||||
"""Computes desired joint positions from a target end-effector pose using inverse kinematics (IK).
|
||||
|
||||
This step translates a Cartesian command (position and orientation of the end-effector) into
|
||||
the corresponding joint-space commands for each motor.
|
||||
|
||||
Attributes:
|
||||
kinematics: The robot's kinematic model for inverse kinematics.
|
||||
motor_names: A list of motor names for which to compute joint positions.
|
||||
q_curr: Internal state storing the last joint positions, used as an initial guess for the IK solver.
|
||||
initial_guess_current_joints: If True, use the robot's current joint state as the IK guess.
|
||||
If False, use the solution from the previous step.
|
||||
orientation_weight: Weight for the orientation constraint passed to
|
||||
``RobotKinematics.inverse_kinematics``. Defaults to ``0.01`` (matching the solver
|
||||
default, so existing callers are unchanged). Set to ``0.0`` for position-only IK on
|
||||
under-actuated arms; a small nonzero weight gives soft-orientation IK on the 5-DOF
|
||||
SO-101, where the wrist tracks orientation only partially (position dominates).
|
||||
**Attributes**:
|
||||
- **kinematics** (`RobotKinematics`) -- The robot's kinematic model for inverse kinematics.
|
||||
- **motor_names** (`list[str]`) -- A list of motor names for which to compute joint positions.
|
||||
- **q_curr** (`np.ndarray | None`) -- Internal state storing the last joint positions, used as an
|
||||
initial guess for the IK solver.
|
||||
- **initial_guess_current_joints** (`bool`) -- If True, use the robot's current joint state as the IK
|
||||
guess. If False, use the solution from the previous step.
|
||||
- **orientation_weight** (`float`) -- Weight for the orientation constraint passed to
|
||||
``RobotKinematics.inverse_kinematics``. Defaults to ``0.01`` (matching the solver default, so
|
||||
existing callers are unchanged). Set to ``0.0`` for position-only IK on under-actuated arms; a small
|
||||
nonzero weight gives soft-orientation IK on the 5-DOF SO-101, where the wrist tracks orientation
|
||||
only partially (position dominates).
|
||||
"""
|
||||
|
||||
kinematics: RobotKinematics
|
||||
@@ -300,6 +336,15 @@ class InverseKinematicsEEToJoints(RobotActionProcessorStep):
|
||||
orientation_weight: float = 0.01
|
||||
|
||||
def action(self, action: RobotAction) -> RobotAction:
|
||||
"""Transform the action for this step.
|
||||
|
||||
Args:
|
||||
action (`dict[str, Any]`):
|
||||
The incoming robot action.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The transformed action.
|
||||
"""
|
||||
x = action.pop("ee.x")
|
||||
y = action.pop("ee.y")
|
||||
z = action.pop("ee.z")
|
||||
@@ -355,6 +400,16 @@ class InverseKinematicsEEToJoints(RobotActionProcessorStep):
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
"""Update the feature contract to match what this step does to the data.
|
||||
|
||||
Args:
|
||||
features (`dict[PipelineFeatureType, dict[str, PolicyFeature]]`):
|
||||
The pipeline's feature contract so far.
|
||||
|
||||
Returns:
|
||||
`dict[PipelineFeatureType, dict[str, PolicyFeature]]`: The contract with this step's key changes
|
||||
applied.
|
||||
"""
|
||||
for feat in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
|
||||
features[PipelineFeatureType.ACTION].pop(f"ee.{feat}", None)
|
||||
|
||||
@@ -373,20 +428,20 @@ class InverseKinematicsEEToJoints(RobotActionProcessorStep):
|
||||
@ProcessorStepRegistry.register("gripper_velocity_to_joint")
|
||||
@dataclass
|
||||
class GripperVelocityToJoint(RobotActionProcessorStep):
|
||||
"""
|
||||
Converts a gripper velocity command into a target gripper joint position.
|
||||
"""Converts a gripper velocity command into a target gripper joint position.
|
||||
|
||||
This step integrates a normalized velocity command over time to produce a position command,
|
||||
taking the current gripper position as a starting point. It also supports a discrete mode
|
||||
where integer actions map to open, close, or no-op.
|
||||
|
||||
Attributes:
|
||||
motor_names: A list of motor names, which must include 'gripper'.
|
||||
speed_factor: A scaling factor to convert the normalized velocity command to a position change.
|
||||
clip_min: The minimum allowed gripper joint position.
|
||||
clip_max: The maximum allowed gripper joint position.
|
||||
discrete_gripper: If True, interpret the input as a discrete class index
|
||||
{0 = close, 1 = stay, 2 = open}, matching `GamepadTeleop.GripperAction`.
|
||||
**Attributes**:
|
||||
- **motor_names** -- A list of motor names, which must include 'gripper'.
|
||||
- **speed_factor** (`float`) -- A scaling factor to convert the normalized velocity command to a
|
||||
position change.
|
||||
- **clip_min** (`float`) -- The minimum allowed gripper joint position.
|
||||
- **clip_max** (`float`) -- The maximum allowed gripper joint position.
|
||||
- **discrete_gripper** (`bool`) -- If True, interpret the input as a discrete class index {0 = close,
|
||||
1 = stay, 2 = open}, matching `GamepadTeleop.GripperAction`.
|
||||
"""
|
||||
|
||||
speed_factor: float = 20.0
|
||||
@@ -395,6 +450,15 @@ class GripperVelocityToJoint(RobotActionProcessorStep):
|
||||
discrete_gripper: bool = False
|
||||
|
||||
def action(self, action: RobotAction) -> RobotAction:
|
||||
"""Transform the action for this step.
|
||||
|
||||
Args:
|
||||
action (`dict[str, Any]`):
|
||||
The incoming robot action.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The transformed action.
|
||||
"""
|
||||
raw_observation = self.transition.get(TransitionKey.OBSERVATION)
|
||||
|
||||
gripper_vel = action.pop("ee.gripper_vel")
|
||||
@@ -428,6 +492,16 @@ class GripperVelocityToJoint(RobotActionProcessorStep):
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
"""Update the feature contract to match what this step does to the data.
|
||||
|
||||
Args:
|
||||
features (`dict[PipelineFeatureType, dict[str, PolicyFeature]]`):
|
||||
The pipeline's feature contract so far.
|
||||
|
||||
Returns:
|
||||
`dict[PipelineFeatureType, dict[str, PolicyFeature]]`: The contract with this step's key changes
|
||||
applied.
|
||||
"""
|
||||
features[PipelineFeatureType.ACTION].pop("ee.gripper_vel", None)
|
||||
features[PipelineFeatureType.ACTION]["ee.gripper_pos"] = PolicyFeature(
|
||||
type=FeatureType.ACTION, shape=(1,)
|
||||
@@ -439,6 +513,21 @@ class GripperVelocityToJoint(RobotActionProcessorStep):
|
||||
def compute_forward_kinematics_joints_to_ee(
|
||||
joints: dict[str, Any], kinematics: RobotKinematics, motor_names: list[str]
|
||||
) -> dict[str, Any]:
|
||||
"""Replace joint positions with the end-effector pose they produce.
|
||||
|
||||
Args:
|
||||
joints (`dict[str, Any]`):
|
||||
Joint values keyed `"<motor>.pos"`, including `"gripper.pos"`. Modified in place: the joint
|
||||
keys named in `motor_names` are removed.
|
||||
kinematics (`RobotKinematics`):
|
||||
The arm's kinematic model.
|
||||
motor_names (`list[str]`):
|
||||
The motors, in the order the kinematic model expects them.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The same dict with `ee.x`, `ee.y`, `ee.z` for position, `ee.wx`, `ee.wy`,
|
||||
`ee.wz` for orientation as a rotation vector, and `ee.gripper_pos` carried through unchanged.
|
||||
"""
|
||||
motor_joint_values = [joints[f"{n}.pos"] for n in motor_names]
|
||||
|
||||
q = np.array(motor_joint_values, dtype=float)
|
||||
@@ -461,26 +550,44 @@ def compute_forward_kinematics_joints_to_ee(
|
||||
@ProcessorStepRegistry.register("forward_kinematics_joints_to_ee_observation")
|
||||
@dataclass
|
||||
class ForwardKinematicsJointsToEEObservation(ObservationProcessorStep):
|
||||
"""
|
||||
Computes the end-effector pose from joint positions using forward kinematics (FK).
|
||||
"""Computes the end-effector pose from joint positions using forward kinematics (FK).
|
||||
|
||||
This step is typically used to add the robot's Cartesian pose to the observation space,
|
||||
which can be useful for visualization or as an input to a policy.
|
||||
|
||||
Attributes:
|
||||
kinematics: The robot's kinematic model.
|
||||
**Attributes**:
|
||||
- **kinematics** (`RobotKinematics`) -- The robot's kinematic model.
|
||||
"""
|
||||
|
||||
kinematics: RobotKinematics
|
||||
motor_names: list[str]
|
||||
|
||||
def observation(self, observation: RobotObservation) -> RobotObservation:
|
||||
"""Replace the observation's joint positions with the end-effector pose.
|
||||
|
||||
Args:
|
||||
observation (`dict[str, Any]`):
|
||||
The incoming observation, containing `"<motor>.pos"` keys.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The observation with `ee.*` keys in place of the joint positions.
|
||||
"""
|
||||
return compute_forward_kinematics_joints_to_ee(observation, self.kinematics, self.motor_names)
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
# We only use the ee pose in the dataset, so we don't need the joint positions
|
||||
"""Update the feature contract to match what this step does to the data.
|
||||
|
||||
Args:
|
||||
features (`dict[PipelineFeatureType, dict[str, PolicyFeature]]`):
|
||||
The pipeline's feature contract so far.
|
||||
|
||||
Returns:
|
||||
`dict[PipelineFeatureType, dict[str, PolicyFeature]]`: The contract with this step's key changes
|
||||
applied.
|
||||
"""
|
||||
for n in self.motor_names:
|
||||
features[PipelineFeatureType.OBSERVATION].pop(f"{n}.pos", None)
|
||||
# We specify the dataset features of this step that we want to be stored in the dataset
|
||||
@@ -494,26 +601,44 @@ class ForwardKinematicsJointsToEEObservation(ObservationProcessorStep):
|
||||
@ProcessorStepRegistry.register("forward_kinematics_joints_to_ee_action")
|
||||
@dataclass
|
||||
class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
|
||||
"""
|
||||
Computes the end-effector pose from joint positions using forward kinematics (FK).
|
||||
"""Computes the end-effector pose from joint positions using forward kinematics (FK).
|
||||
|
||||
This step is typically used to add the robot's Cartesian pose to the observation space,
|
||||
which can be useful for visualization or as an input to a policy.
|
||||
|
||||
Attributes:
|
||||
kinematics: The robot's kinematic model.
|
||||
**Attributes**:
|
||||
- **kinematics** (`RobotKinematics`) -- The robot's kinematic model.
|
||||
"""
|
||||
|
||||
kinematics: RobotKinematics
|
||||
motor_names: list[str]
|
||||
|
||||
def action(self, action: RobotAction) -> RobotAction:
|
||||
"""Transform the action for this step.
|
||||
|
||||
Args:
|
||||
action (`dict[str, Any]`):
|
||||
The incoming robot action.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The transformed action.
|
||||
"""
|
||||
return compute_forward_kinematics_joints_to_ee(action, self.kinematics, self.motor_names)
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
# We only use the ee pose in the dataset, so we don't need the joint positions
|
||||
"""Update the feature contract to match what this step does to the data.
|
||||
|
||||
Args:
|
||||
features (`dict[PipelineFeatureType, dict[str, PolicyFeature]]`):
|
||||
The pipeline's feature contract so far.
|
||||
|
||||
Returns:
|
||||
`dict[PipelineFeatureType, dict[str, PolicyFeature]]`: The contract with this step's key changes
|
||||
applied.
|
||||
"""
|
||||
for n in self.motor_names:
|
||||
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
|
||||
# Store end-effector features as actions in the dataset schema
|
||||
@@ -527,10 +652,21 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
|
||||
@ProcessorStepRegistry.register(name="forward_kinematics_joints_to_ee")
|
||||
@dataclass
|
||||
class ForwardKinematicsJointsToEE(ProcessorStep):
|
||||
"""Applies forward kinematics to whichever of the action and observation are present.
|
||||
|
||||
A convenience wrapper over [`ForwardKinematicsJointsToEEAction`] and
|
||||
[`ForwardKinematicsJointsToEEObservation`], so a pipeline needs one step instead of two.
|
||||
|
||||
**Attributes**:
|
||||
- **kinematics** (`RobotKinematics`) -- The arm's kinematic model.
|
||||
- **motor_names** (`list[str]`) -- The motors, in the order the kinematic model expects them.
|
||||
"""
|
||||
|
||||
kinematics: RobotKinematics
|
||||
motor_names: list[str]
|
||||
|
||||
def __post_init__(self):
|
||||
"""Build the action and observation sub-steps this step delegates to."""
|
||||
self.joints_to_ee_action_processor = ForwardKinematicsJointsToEEAction(
|
||||
kinematics=self.kinematics, motor_names=self.motor_names
|
||||
)
|
||||
@@ -539,6 +675,15 @@ class ForwardKinematicsJointsToEE(ProcessorStep):
|
||||
)
|
||||
|
||||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||
"""Apply forward kinematics to whichever of the action and observation are present.
|
||||
|
||||
Args:
|
||||
transition (`EnvTransition`):
|
||||
The transition to transform.
|
||||
|
||||
Returns:
|
||||
`EnvTransition`: The transition with `ee.*` keys in place of joint positions.
|
||||
"""
|
||||
if transition.get(TransitionKey.ACTION) is not None:
|
||||
transition = self.joints_to_ee_action_processor(transition)
|
||||
if transition.get(TransitionKey.OBSERVATION) is not None:
|
||||
@@ -548,6 +693,16 @@ class ForwardKinematicsJointsToEE(ProcessorStep):
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
"""Update the feature contract to match what this step does to the data.
|
||||
|
||||
Args:
|
||||
features (`dict[PipelineFeatureType, dict[str, PolicyFeature]]`):
|
||||
The pipeline's feature contract so far.
|
||||
|
||||
Returns:
|
||||
`dict[PipelineFeatureType, dict[str, PolicyFeature]]`: The contract with this step's key changes
|
||||
applied.
|
||||
"""
|
||||
if features[PipelineFeatureType.ACTION] is not None:
|
||||
features = self.joints_to_ee_action_processor.transform_features(features)
|
||||
if features[PipelineFeatureType.OBSERVATION] is not None:
|
||||
@@ -558,8 +713,7 @@ class ForwardKinematicsJointsToEE(ProcessorStep):
|
||||
@ProcessorStepRegistry.register("inverse_kinematics_rl_step")
|
||||
@dataclass
|
||||
class InverseKinematicsRLStep(ProcessorStep):
|
||||
"""
|
||||
Computes desired joint positions from a target end-effector pose using inverse kinematics (IK).
|
||||
"""Computes desired joint positions from a target end-effector pose using inverse kinematics (IK).
|
||||
|
||||
This is modified from the InverseKinematicsEEToJoints step to be used in the RL pipeline.
|
||||
"""
|
||||
@@ -570,6 +724,15 @@ class InverseKinematicsRLStep(ProcessorStep):
|
||||
initial_guess_current_joints: bool = True
|
||||
|
||||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||
"""Solve inverse kinematics for the transition's end-effector action.
|
||||
|
||||
Args:
|
||||
transition (`EnvTransition`):
|
||||
The transition to transform.
|
||||
|
||||
Returns:
|
||||
`EnvTransition`: The transition with joint targets in place of the `ee.*` action.
|
||||
"""
|
||||
new_transition = dict(transition)
|
||||
action = new_transition.get(TransitionKey.ACTION)
|
||||
if action is None:
|
||||
@@ -633,6 +796,16 @@ class InverseKinematicsRLStep(ProcessorStep):
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
"""Update the feature contract to match what this step does to the data.
|
||||
|
||||
Args:
|
||||
features (`dict[PipelineFeatureType, dict[str, PolicyFeature]]`):
|
||||
The pipeline's feature contract so far.
|
||||
|
||||
Returns:
|
||||
`dict[PipelineFeatureType, dict[str, PolicyFeature]]`: The contract with this step's key changes
|
||||
applied.
|
||||
"""
|
||||
for feat in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
|
||||
features[PipelineFeatureType.ACTION].pop(f"ee.{feat}", None)
|
||||
|
||||
|
||||
@@ -35,15 +35,35 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SOFollower(Robot):
|
||||
"""
|
||||
Generic SO follower base implementing common functionality for SO-100/101/10X.
|
||||
Designed to be subclassed with a per-hardware-model `config_class` and `name`.
|
||||
"""The SO-family follower arm: a 5-DOF arm plus gripper on a Feetech bus.
|
||||
|
||||
`SO100Follower` and `SO101Follower` are aliases of this class. The two arms differ in calibration and
|
||||
gearing, not control code, so both are driven through the same implementation with a different
|
||||
`config_class` and `name`.
|
||||
|
||||
Actions and observations are keyed `"<motor>.pos"`; cameras named in the config appear in observations
|
||||
under their own keys. See [`~robots.Robot`] for the contract every method here implements.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> from lerobot.robots.so_follower import SO101Follower, SO101FollowerConfig
|
||||
>>> robot = SO101Follower(SO101FollowerConfig(port="/dev/ttyACM0")) # doctest: +SKIP
|
||||
>>> with robot: # doctest: +SKIP
|
||||
... observation = robot.get_observation()
|
||||
... robot.send_action({"shoulder_pan.pos": 0.0})
|
||||
```
|
||||
"""
|
||||
|
||||
config_class = SOFollowerRobotConfig
|
||||
name = "so_follower"
|
||||
|
||||
def __init__(self, config: SOFollowerRobotConfig):
|
||||
"""Build the robot from its configuration.
|
||||
|
||||
Args:
|
||||
config (`SOFollowerRobotConfig`):
|
||||
The robot's configuration. Its `port` and `cameras` determine what is connected.
|
||||
"""
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
# choose normalization mode depending on config if available
|
||||
@@ -78,23 +98,48 @@ class SOFollower(Robot):
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
"""The arm's joint positions plus one entry per configured camera.
|
||||
|
||||
Returns:
|
||||
`dict[str, type | tuple]`: `"<motor>.pos"` keys mapped to `float`, and one key per camera
|
||||
mapped to its `(height, width, channels)` shape.
|
||||
"""
|
||||
return {**self._motors_ft, **self._cameras_ft}
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
"""The arm's goal joint positions.
|
||||
|
||||
Returns:
|
||||
`dict[str, type]`: `"<motor>.pos"` keys mapped to `float`.
|
||||
"""
|
||||
return self._motors_ft
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Whether the motor bus and every configured camera are connected.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` only when all of them are.
|
||||
"""
|
||||
return self.bus.is_connected and all(cam.is_connected for cam in self.cameras.values())
|
||||
|
||||
@check_if_already_connected
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""
|
||||
We assume that at connection time, arm is in a rest position,
|
||||
and torque can be safely disabled to run calibration.
|
||||
"""
|
||||
"""Connect the motor bus and cameras, calibrating and configuring the arm.
|
||||
|
||||
> [!WARNING]
|
||||
> The arm is assumed to be at rest when this is called, because torque is disabled to run
|
||||
> calibration. Do not call it with the arm holding a load.
|
||||
|
||||
Args:
|
||||
calibrate (`bool`, *optional*, defaults to `True`):
|
||||
Whether to run calibration when the motors disagree with the calibration file, or no file
|
||||
exists yet. Calibration is interactive and prompts on stdin.
|
||||
|
||||
Raises:
|
||||
DeviceAlreadyConnectedError: If the robot is already connected.
|
||||
"""
|
||||
self.bus.connect()
|
||||
if not self.is_calibrated and calibrate:
|
||||
logger.info(
|
||||
@@ -110,9 +155,19 @@ class SOFollower(Robot):
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
"""Whether the motors' stored calibration matches the calibration file.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` when the arm needs no recalibration.
|
||||
"""
|
||||
return self.bus.is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
"""Calibrate the arm, writing the result to the motors and the calibration file.
|
||||
|
||||
This is interactive: it prompts on stdin to reuse an existing calibration file, and otherwise asks
|
||||
you to move the arm to its middle position and then through each joint's full range.
|
||||
"""
|
||||
if self.calibration:
|
||||
# Calibration file exists, ask user whether to use it or run new calibration
|
||||
user_input = input(
|
||||
@@ -157,6 +212,11 @@ class SOFollower(Robot):
|
||||
print("Calibration saved to", self.calibration_fpath)
|
||||
|
||||
def configure(self) -> None:
|
||||
"""Write the position-mode operating mode and the configured PID gains to every motor.
|
||||
|
||||
The gripper additionally gets reduced torque, current and overload limits so that gripping a rigid
|
||||
object does not burn out its motor.
|
||||
"""
|
||||
with self.bus.torque_disabled():
|
||||
self.bus.configure_motors()
|
||||
for motor in self.bus.motors:
|
||||
@@ -171,6 +231,11 @@ class SOFollower(Robot):
|
||||
self.bus.write("Overload_Torque", motor, 25) # 25% torque when overloaded
|
||||
|
||||
def setup_motors(self) -> None:
|
||||
"""Assign each motor its bus ID, one at a time.
|
||||
|
||||
Run this once when building an arm. It is interactive: it prompts you to connect the controller
|
||||
board to a single motor at a time, working from the gripper back to the base.
|
||||
"""
|
||||
for motor in reversed(self.bus.motors):
|
||||
input(f"Connect the controller board to the '{motor}' motor only and press enter.")
|
||||
self.bus.setup_motor(motor)
|
||||
@@ -178,6 +243,14 @@ class SOFollower(Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
"""Read the arm's joint positions and one frame from each camera.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: Keys matching [`~robots.Robot.observation_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
# Read arm position
|
||||
start = time.perf_counter()
|
||||
obs_dict = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
|
||||
@@ -215,7 +288,6 @@ class SOFollower(Robot):
|
||||
Returns:
|
||||
RobotAction: the action sent to the motors, potentially clipped.
|
||||
"""
|
||||
|
||||
goal_pos = {key.removesuffix(".pos"): val for key, val in action.items() if key.endswith(".pos")}
|
||||
|
||||
# Cap goal position when too far away from present position.
|
||||
@@ -231,6 +303,13 @@ class SOFollower(Robot):
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self):
|
||||
"""Disconnect the motor bus and every camera.
|
||||
|
||||
Torque is released first unless `disable_torque_on_disconnect` is `False`.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
self.bus.disconnect(self.config.disable_torque_on_disconnect)
|
||||
for cam in self.cameras.values():
|
||||
cam.disconnect()
|
||||
|
||||
@@ -47,6 +47,42 @@ _DEFAULT_KP, _DEFAULT_KD = _build_gains()
|
||||
@RobotConfig.register_subclass("unitree_g1")
|
||||
@dataclass
|
||||
class UnitreeG1Config(RobotConfig):
|
||||
"""Configuration for the Unitree G1 humanoid.
|
||||
|
||||
The G1 is reached over a ZMQ bridge rather than a serial bus, so there is no `port` field and
|
||||
calibration is handled by the robot's own firmware.
|
||||
|
||||
All 29 joints are addressed by index, so `kp`, `kd` and `default_positions` are lists in the G1's joint
|
||||
order: left leg, right leg, waist, left arm, left wrist, right arm, right wrist.
|
||||
|
||||
Args:
|
||||
kp (`list[float]`, *optional*):
|
||||
Per-joint proportional gains, 29 values. Defaults to the per-body-part gains recommended by
|
||||
Unitree.
|
||||
kd (`list[float]`, *optional*):
|
||||
Per-joint derivative gains, 29 values.
|
||||
default_positions (`list[float]`, *optional*):
|
||||
Per-joint home positions, 29 values. Defaults to all zeros.
|
||||
control_dt (`float`, *optional*, defaults to 0.004):
|
||||
Control loop timestep in seconds, i.e. 250 Hz.
|
||||
is_simulation (`bool`, *optional*, defaults to `True`):
|
||||
Whether to drive a MuJoCo simulation instead of the physical robot. Keep `True` until the
|
||||
behaviour is validated in sim.
|
||||
robot_ip (`str`, *optional*, defaults to `"192.168.123.164"`):
|
||||
Address of the robot's ZMQ bridge. The default is the G1's factory address.
|
||||
cameras (`dict[str, CameraConfig]`, *optional*):
|
||||
ZMQ-based remote cameras to read alongside the joint states.
|
||||
gravity_compensation (`bool`, *optional*, defaults to `False`):
|
||||
Whether to compensate for gravity on the arms using the arm IK solver.
|
||||
controller (`str`, *optional*):
|
||||
Class name of the lower-body locomotion controller, e.g. `"GrootLocomotionController"` or
|
||||
`"HolosomaLocomotionController"`. `None` leaves the legs uncontrolled.
|
||||
id (`str`, *optional*):
|
||||
Identifier for this particular robot.
|
||||
calibration_dir (`Path`, *optional*):
|
||||
Unused: the G1 manages its own calibration.
|
||||
"""
|
||||
|
||||
kp: list[float] = field(default_factory=lambda: _DEFAULT_KP.copy())
|
||||
kd: list[float] = field(default_factory=lambda: _DEFAULT_KD.copy())
|
||||
|
||||
|
||||
@@ -24,7 +24,17 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WeightedMovingFilter:
|
||||
"""A fixed-length weighted moving average over recent samples, used to smooth IK solutions."""
|
||||
|
||||
def __init__(self, weights, data_size=14):
|
||||
"""Set up the filter.
|
||||
|
||||
Args:
|
||||
weights:
|
||||
Per-sample weights, newest first. Their length sets the window size.
|
||||
data_size (`int`, *optional*, defaults to 14):
|
||||
Number of values in each sample.
|
||||
"""
|
||||
self._window_size = len(weights)
|
||||
self._weights = np.array(weights)
|
||||
self._data_size = data_size
|
||||
@@ -39,6 +49,12 @@ class WeightedMovingFilter:
|
||||
return data_array.T @ self._weights
|
||||
|
||||
def add_data(self, new_data):
|
||||
"""Push a sample into the window and recompute the filtered value.
|
||||
|
||||
Args:
|
||||
new_data:
|
||||
A sample of length `data_size`. Ignored if identical to the newest one already held.
|
||||
"""
|
||||
assert len(new_data) == self._data_size
|
||||
|
||||
if len(self._data_queue) > 0 and np.array_equal(
|
||||
@@ -51,11 +67,24 @@ class WeightedMovingFilter:
|
||||
|
||||
@property
|
||||
def filtered_data(self):
|
||||
"""The current weighted average.
|
||||
|
||||
Returns:
|
||||
`np.ndarray`: The filtered sample.
|
||||
"""
|
||||
return self._filtered_data
|
||||
|
||||
|
||||
class G1_29_ArmIK: # noqa: N801
|
||||
"""Inverse kinematics for the G1's two arms, solved together as one optimisation problem."""
|
||||
|
||||
def __init__(self, unit_test=False):
|
||||
"""Build the arm model and the IK solver.
|
||||
|
||||
Args:
|
||||
unit_test (`bool`, *optional*, defaults to `False`):
|
||||
Whether to run in test mode, which visualises the solution instead of driving a robot.
|
||||
"""
|
||||
import casadi
|
||||
import pinocchio as pin
|
||||
from huggingface_hub import snapshot_download
|
||||
@@ -230,6 +259,21 @@ class G1_29_ArmIK: # noqa: N801
|
||||
self.smooth_filter = WeightedMovingFilter(np.array([0.4, 0.3, 0.2, 0.1]), 14)
|
||||
|
||||
def solve_ik(self, left_wrist, right_wrist, current_lr_arm_motor_q=None, current_lr_arm_motor_dq=None):
|
||||
"""Solve for the arm joint angles that place both wrists at the requested poses.
|
||||
|
||||
Args:
|
||||
left_wrist:
|
||||
Target pose of the left wrist as a 4x4 homogeneous transform.
|
||||
right_wrist:
|
||||
Target pose of the right wrist as a 4x4 homogeneous transform.
|
||||
current_lr_arm_motor_q (*optional*):
|
||||
Present arm joint positions, used as the solver's initial guess.
|
||||
current_lr_arm_motor_dq (*optional*):
|
||||
Present arm joint velocities, used to compute feed-forward torques.
|
||||
|
||||
Returns:
|
||||
`tuple`: The solved joint positions and the corresponding torques.
|
||||
"""
|
||||
if current_lr_arm_motor_q is not None:
|
||||
self.init_data = current_lr_arm_motor_q
|
||||
self.opti.set_initial(self.var_q, self.init_data)
|
||||
@@ -268,6 +312,17 @@ class G1_29_ArmIK: # noqa: N801
|
||||
return sol_q, sol_tauff
|
||||
|
||||
def solve_tau(self, current_lr_arm_motor_q=None, current_lr_arm_motor_dq=None):
|
||||
"""Compute the gravity-compensating torques for the arms at a given state.
|
||||
|
||||
Args:
|
||||
current_lr_arm_motor_q (*optional*):
|
||||
Present arm joint positions.
|
||||
current_lr_arm_motor_dq (*optional*):
|
||||
Present arm joint velocities.
|
||||
|
||||
Returns:
|
||||
`np.ndarray`: Per-joint torques.
|
||||
"""
|
||||
try:
|
||||
q_g1 = np.array(current_lr_arm_motor_q, dtype=float)
|
||||
if q_g1.shape[0] != len(self._arm_joint_names_g1):
|
||||
|
||||
@@ -44,6 +44,8 @@ def get_gravity_orientation(quaternion: list[float] | np.ndarray) -> np.ndarray:
|
||||
|
||||
|
||||
class G1_29_JointArmIndex(IntEnum):
|
||||
"""Indices of the G1's arm and wrist joints within its 29-joint state vector."""
|
||||
|
||||
# Left arm
|
||||
kLeftShoulderPitch = 15
|
||||
kLeftShoulderRoll = 16
|
||||
@@ -79,6 +81,8 @@ def make_locomotion_controller(name: str | None):
|
||||
|
||||
|
||||
class G1_29_JointIndex(IntEnum):
|
||||
"""Indices of all 29 G1 joints, in the order the robot reports and accepts them."""
|
||||
|
||||
# Left leg
|
||||
kLeftHipPitch = 0
|
||||
kLeftHipRoll = 1
|
||||
|
||||
@@ -83,6 +83,7 @@ class GrootLocomotionController:
|
||||
control_dt = CONTROL_DT # Expose for unitree_g1.py
|
||||
|
||||
def __init__(self):
|
||||
"""Load the GR00T locomotion policy and set up its observation history."""
|
||||
# Load policies
|
||||
self.policy_balance, self.policy_walk = load_groot_policies()
|
||||
|
||||
|
||||
@@ -101,6 +101,7 @@ class HolosomaLocomotionController:
|
||||
control_dt = CONTROL_DT # Expose for unitree_g1.py
|
||||
|
||||
def __init__(self):
|
||||
"""Load the HoloSoma locomotion policy and set up its observation history."""
|
||||
# Load policy and gains
|
||||
self.policy, self.kp, self.kd = load_policy()
|
||||
|
||||
|
||||
@@ -14,8 +14,7 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
DDS-to-ZMQ bridge server for Unitree G1 robot.
|
||||
"""DDS-to-ZMQ bridge server for Unitree G1 robot.
|
||||
|
||||
This server runs on the robot and forwards:
|
||||
- Robot state (LowState) from DDS to ZMQ (for remote clients)
|
||||
|
||||
@@ -67,11 +67,27 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
@runtime_checkable
|
||||
class LocomotionController(Protocol):
|
||||
"""The interface a lower-body locomotion controller must provide to drive the G1's legs."""
|
||||
|
||||
control_dt: float
|
||||
|
||||
def run_step(self, action: dict, lowstate) -> dict: ...
|
||||
def run_step(self, action: dict, lowstate) -> dict:
|
||||
"""Compute one step of leg commands.
|
||||
|
||||
def reset(self) -> None: ...
|
||||
Args:
|
||||
action (`dict`):
|
||||
The upper-body action and locomotion command for this step.
|
||||
lowstate:
|
||||
The robot's most recent low-level state.
|
||||
|
||||
Returns:
|
||||
`dict`: Leg joint targets for this step.
|
||||
"""
|
||||
...
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Clear any internal state, e.g. an observation history, before a new episode."""
|
||||
...
|
||||
|
||||
|
||||
# DDS topic names follow Unitree SDK naming conventions
|
||||
@@ -82,6 +98,8 @@ kTopicLowState = "rt/lowstate"
|
||||
|
||||
@dataclass
|
||||
class MotorState:
|
||||
"""One motor's reported position, velocity and torque."""
|
||||
|
||||
q: float | None = None # position
|
||||
dq: float | None = None # velocity
|
||||
tau_est: float | None = None # estimated torque
|
||||
@@ -90,6 +108,8 @@ class MotorState:
|
||||
|
||||
@dataclass
|
||||
class IMUState:
|
||||
"""The G1's inertial measurements: orientation, angular velocity and acceleration."""
|
||||
|
||||
quaternion: np.ndarray | None = None # [w, x, y, z]
|
||||
gyroscope: np.ndarray | None = None # [x, y, z] angular velocity (rad/s)
|
||||
accelerometer: np.ndarray | None = None # [x, y, z] linear acceleration (m/s²)
|
||||
@@ -100,6 +120,8 @@ class IMUState:
|
||||
# g1 observation class
|
||||
@dataclass
|
||||
class G1_29_LowState: # noqa: N801
|
||||
"""A full low-level state frame: every motor's state plus the IMU."""
|
||||
|
||||
motor_state: list[MotorState] = field(default_factory=lambda: [MotorState() for _ in G1_29_JointIndex])
|
||||
imu_state: IMUState = field(default_factory=IMUState)
|
||||
wireless_remote: bytes | None = None # Raw wireless remote data
|
||||
@@ -107,10 +129,29 @@ class G1_29_LowState: # noqa: N801
|
||||
|
||||
|
||||
class UnitreeG1(Robot):
|
||||
"""The Unitree G1 humanoid, driven over a ZMQ bridge.
|
||||
|
||||
Upper-body joints are commanded directly. The legs are handled by an optional locomotion controller
|
||||
named in the config, which runs its own loop against the robot's low-level state. Set
|
||||
`is_simulation=True` to drive a MuJoCo model instead of the physical robot.
|
||||
|
||||
See [`~robots.Robot`] for the contract every method here implements.
|
||||
"""
|
||||
|
||||
config_class = UnitreeG1Config
|
||||
name = "unitree_g1"
|
||||
|
||||
def __init__(self, config: UnitreeG1Config):
|
||||
"""Build the robot and, if one is configured, its locomotion controller.
|
||||
|
||||
Args:
|
||||
config (`UnitreeG1Config`):
|
||||
The robot's configuration, including gains, the ZMQ bridge address and whether to run
|
||||
against MuJoCo instead of hardware.
|
||||
|
||||
Raises:
|
||||
ImportError: If the `unitree_g1` extra is not installed.
|
||||
"""
|
||||
require_package("unitree-sdk2py", extra="unitree_g1", import_name="unitree_sdk2py")
|
||||
super().__init__(config)
|
||||
|
||||
@@ -204,6 +245,18 @@ class UnitreeG1(Robot):
|
||||
kd: np.ndarray | list[float] | None = None,
|
||||
tau: np.ndarray | list[float] | None = None,
|
||||
) -> None: # writes robot command whenever requested
|
||||
"""Write a low-level command frame to the robot.
|
||||
|
||||
Args:
|
||||
action (`dict[str, Any]`):
|
||||
Target joint positions for this step.
|
||||
kp (`np.ndarray | list[float]`, *optional*):
|
||||
Per-joint proportional gains. Defaults to the config's `kp`.
|
||||
kd (`np.ndarray | list[float]`, *optional*):
|
||||
Per-joint derivative gains. Defaults to the config's `kd`.
|
||||
tau (`np.ndarray | list[float]`, *optional*):
|
||||
Per-joint feed-forward torques. Defaults to zero.
|
||||
"""
|
||||
for motor in G1_29_JointIndex:
|
||||
key = f"{motor.name}.q"
|
||||
if key in action:
|
||||
@@ -233,10 +286,21 @@ class UnitreeG1(Robot):
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
"""The values this robot reports, and their types or shapes.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys as returned by [`~robots.Robot.get_observation`], mapped to a scalar type for
|
||||
proprioceptive values or to a `(height, width, channels)` shape for images.
|
||||
"""
|
||||
return {**self._motors_ft, **self._cameras_ft}
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
"""The values this robot accepts, and their types.
|
||||
|
||||
Returns:
|
||||
`dict`: Keys accepted by [`~robots.Robot.send_action`], mapped to their type.
|
||||
"""
|
||||
if self.controller is None:
|
||||
return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
|
||||
|
||||
@@ -288,13 +352,27 @@ class UnitreeG1(Robot):
|
||||
|
||||
def calibrate(self) -> None:
|
||||
# TODO: implement g1_29 calibration
|
||||
"""Calibrate the robot and store the result.
|
||||
|
||||
Interactive: prompts on stdin and asks you to move the robot through the required positions.
|
||||
"""
|
||||
pass
|
||||
|
||||
def configure(self) -> None:
|
||||
"""Apply the operating mode, gains and limits from the configuration to the robot."""
|
||||
pass
|
||||
|
||||
def connect(self, calibrate: bool = True) -> None: # connect to DDS
|
||||
# Initialize DDS channel and simulation environment
|
||||
"""Connect to the robot and its cameras, then apply the configured settings.
|
||||
|
||||
Args:
|
||||
calibrate (`bool`, *optional*, defaults to `True`):
|
||||
Whether to run calibration if the robot is not already calibrated.
|
||||
|
||||
Raises:
|
||||
DeviceAlreadyConnectedError: If the robot is already connected.
|
||||
"""
|
||||
if self.config.is_simulation:
|
||||
from lerobot.envs import make_env
|
||||
|
||||
@@ -373,6 +451,11 @@ class UnitreeG1(Robot):
|
||||
|
||||
def disconnect(self):
|
||||
# Put robot in passive mode before stopping threads
|
||||
"""Disconnect from the robot and its cameras.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
if not self.config.is_simulation:
|
||||
self._send_zero_torque()
|
||||
|
||||
@@ -417,6 +500,14 @@ class UnitreeG1(Robot):
|
||||
cam.disconnect()
|
||||
|
||||
def get_observation(self) -> RobotObservation:
|
||||
"""Read the robot's current state and a frame from each camera.
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: Keys matching [`~robots.Robot.observation_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
with self._lowstate_lock:
|
||||
lowstate = self._lowstate
|
||||
if lowstate is None:
|
||||
@@ -471,6 +562,18 @@ class UnitreeG1(Robot):
|
||||
return obs
|
||||
|
||||
def send_action(self, action: RobotAction) -> RobotAction:
|
||||
"""Command the robot to move towards a target configuration.
|
||||
|
||||
Args:
|
||||
action (`dict[str, Any]`):
|
||||
Target values, keyed as in [`~robots.Robot.action_features`].
|
||||
|
||||
Returns:
|
||||
`dict[str, Any]`: The action actually sent, which may be clipped by `max_relative_target`.
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If the robot is not connected.
|
||||
"""
|
||||
action_to_publish = action
|
||||
if self.controller is not None:
|
||||
# Controller thread owns legs/waist. Here we only update joystick inputs
|
||||
@@ -511,10 +614,20 @@ class UnitreeG1(Robot):
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
"""Whether the robot is calibrated.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` when no calibration is needed before use.
|
||||
"""
|
||||
return True
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
"""Whether every device this robot uses is connected.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` only when the robot and all its cameras are connected.
|
||||
"""
|
||||
with self._lowstate_lock:
|
||||
return self._lowstate is not None
|
||||
|
||||
@@ -525,6 +638,11 @@ class UnitreeG1(Robot):
|
||||
|
||||
@property
|
||||
def cameras(self) -> dict:
|
||||
"""The robot's configured cameras.
|
||||
|
||||
Returns:
|
||||
`dict`: Camera name mapped to its instance.
|
||||
"""
|
||||
return self._cameras
|
||||
|
||||
def reset(
|
||||
@@ -532,6 +650,18 @@ class UnitreeG1(Robot):
|
||||
control_dt: float | None = None,
|
||||
default_positions: list[float] | None = None,
|
||||
) -> None: # move robot to default position
|
||||
"""Move the robot smoothly to its default joint positions.
|
||||
|
||||
> [!WARNING]
|
||||
> This drives every joint, legs included. Make sure the robot is supported or in a safe posture
|
||||
> before calling it.
|
||||
|
||||
Args:
|
||||
control_dt (`float`, *optional*):
|
||||
Control loop timestep for the move. Defaults to the config's `control_dt`.
|
||||
default_positions (`list[float]`, *optional*):
|
||||
Target positions, 29 values in joint order. Defaults to the config's `default_positions`.
|
||||
"""
|
||||
if control_dt is None:
|
||||
control_dt = self.config.control_dt
|
||||
if default_positions is None:
|
||||
|
||||
@@ -37,8 +37,7 @@ kTopicLowCommand_Debug = "rt/lowcmd"
|
||||
|
||||
|
||||
class LowStateMsg:
|
||||
"""
|
||||
Wrapper class that mimics the Unitree SDK LowState_ message structure.
|
||||
"""Wrapper class that mimics the Unitree SDK LowState_ message structure.
|
||||
|
||||
Reconstructs the message from deserialized JSON data to maintain
|
||||
compatibility with existing code that expects SDK message objects.
|
||||
@@ -48,6 +47,12 @@ class LowStateMsg:
|
||||
"""Motor state data for a single joint."""
|
||||
|
||||
def __init__(self, data: dict[str, Any]) -> None:
|
||||
"""Build one motor's state from a deserialized JSON frame.
|
||||
|
||||
Args:
|
||||
data (`dict[str, Any]`):
|
||||
The motor's entry from the robot's state message.
|
||||
"""
|
||||
self.q: float = data.get("q", 0.0)
|
||||
self.dq: float = data.get("dq", 0.0)
|
||||
self.tau_est: float = data.get("tau_est", 0.0)
|
||||
@@ -57,6 +62,12 @@ class LowStateMsg:
|
||||
"""IMU sensor data."""
|
||||
|
||||
def __init__(self, data: dict[str, Any]) -> None:
|
||||
"""Build the IMU state from a deserialized JSON frame.
|
||||
|
||||
Args:
|
||||
data (`dict[str, Any]`):
|
||||
The IMU's entry from the robot's state message.
|
||||
"""
|
||||
self.quaternion: list[float] = data.get("quaternion", [1.0, 0.0, 0.0, 0.0])
|
||||
self.gyroscope: list[float] = data.get("gyroscope", [0.0, 0.0, 0.0])
|
||||
self.accelerometer: list[float] = data.get("accelerometer", [0.0, 0.0, 0.0])
|
||||
@@ -100,15 +111,16 @@ def lowcmd_to_dict(topic: str, msg: Any) -> dict[str, Any]:
|
||||
|
||||
|
||||
def ChannelFactoryInitialize(domain_id: int = 0, config: Any = None) -> None: # noqa: N802
|
||||
"""
|
||||
Initialize ZMQ sockets for robot communication.
|
||||
"""Initialize ZMQ sockets for robot communication.
|
||||
|
||||
This function mimics the Unitree SDK's ChannelFactoryInitialize but uses
|
||||
ZMQ sockets to connect to the robot server bridge instead of DDS.
|
||||
|
||||
Args:
|
||||
domain_id: Ignored (for API compatibility with Unitree SDK)
|
||||
config: UnitreeG1Config instance with robot_ip
|
||||
domain_id (`int`, *optional*, defaults to 0):
|
||||
Ignored. Accepted only for API compatibility with the Unitree SDK.
|
||||
config (`Any`, *optional*):
|
||||
A `UnitreeG1Config` supplying `robot_ip`. Defaults to a fresh `UnitreeG1Config` when `None`.
|
||||
"""
|
||||
global _ctx, _lowcmd_sock, _lowstate_sock
|
||||
|
||||
@@ -138,6 +150,14 @@ class ChannelPublisher:
|
||||
"""ZMQ-based publisher that sends commands to the robot server."""
|
||||
|
||||
def __init__(self, topic: str, msg_type: type) -> None:
|
||||
"""Bind the publisher to a topic.
|
||||
|
||||
Args:
|
||||
topic (`str`):
|
||||
The topic name to publish under.
|
||||
msg_type (`type`):
|
||||
The message class this topic carries.
|
||||
"""
|
||||
self.topic = topic
|
||||
self.msg_type = msg_type
|
||||
|
||||
@@ -158,6 +178,14 @@ class ChannelSubscriber:
|
||||
"""ZMQ-based subscriber that receives state from the robot server."""
|
||||
|
||||
def __init__(self, topic: str, msg_type: type) -> None:
|
||||
"""Bind the subscriber to a topic.
|
||||
|
||||
Args:
|
||||
topic (`str`):
|
||||
The topic name to receive from.
|
||||
msg_type (`type`):
|
||||
The message class this topic carries.
|
||||
"""
|
||||
self.topic = topic
|
||||
self.msg_type = msg_type
|
||||
|
||||
|
||||
@@ -23,6 +23,21 @@ from .robot import Robot
|
||||
|
||||
|
||||
def make_robot_from_config(config: RobotConfig) -> Robot:
|
||||
"""Build the robot a configuration selects.
|
||||
|
||||
Dispatches on `config.type`, the name the config registered itself under, so `--robot.type=so101_follower`
|
||||
on the command line produces an `SO101Follower` here.
|
||||
|
||||
Args:
|
||||
config (`RobotConfig`):
|
||||
The configuration to build from.
|
||||
|
||||
Returns:
|
||||
`Robot`: A robot of the type `config` selects, not yet connected.
|
||||
|
||||
Raises:
|
||||
ValueError: If `config.type` matches no built-in robot and no device class can be resolved from it.
|
||||
"""
|
||||
# TODO(Steven): Consider just using the make_device_from_device_class for all types
|
||||
if config.type == "koch_follower":
|
||||
from .koch_follower import KochFollower
|
||||
@@ -91,8 +106,35 @@ def make_robot_from_config(config: RobotConfig) -> Robot:
|
||||
def ensure_safe_goal_position(
|
||||
goal_present_pos: dict[str, tuple[float, float]], max_relative_target: float | dict[str, float]
|
||||
) -> dict[str, float]:
|
||||
"""Caps relative action target magnitude for safety."""
|
||||
"""Cap the magnitude of a relative action target for safety.
|
||||
|
||||
Used by [`~robots.Robot.send_action`] implementations to stop a large jump between the present and goal
|
||||
positions from being written straight to the motors.
|
||||
|
||||
Args:
|
||||
goal_present_pos (`dict[str, tuple[float, float]]`):
|
||||
Maps motor name to a `(goal_position, present_position)` pair.
|
||||
max_relative_target (`float | dict[str, float]`):
|
||||
The largest change from the present position that may be applied. A float caps every motor
|
||||
equally; a dict gives a per-motor cap and must have exactly the same keys as
|
||||
`goal_present_pos`.
|
||||
|
||||
Returns:
|
||||
`dict[str, float]`: The capped goal position for each motor.
|
||||
|
||||
Raises:
|
||||
ValueError: If `max_relative_target` is a dict whose keys differ from those of `goal_present_pos`.
|
||||
TypeError: If `max_relative_target` is neither a float nor a dict.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> from lerobot.robots.utils import ensure_safe_goal_position
|
||||
>>> ensure_safe_goal_position({"shoulder_pan": (50.0, 10.0)}, 5.0)
|
||||
{'shoulder_pan': 15.0}
|
||||
>>> ensure_safe_goal_position({"shoulder_pan": (12.0, 10.0)}, 5.0)
|
||||
{'shoulder_pan': 12.0}
|
||||
```
|
||||
"""
|
||||
if isinstance(max_relative_target, float):
|
||||
diff_cap = dict.fromkeys(goal_present_pos, max_relative_target)
|
||||
elif isinstance(max_relative_target, dict):
|
||||
|
||||
@@ -25,11 +25,6 @@ quantile statistics (q01, q10, q50, q90, q99) in their metadata. This script:
|
||||
3. If missing, computes quantile statistics for all features
|
||||
4. Updates the dataset metadata with the new quantile statistics
|
||||
|
||||
Statistics are accumulated into a single running histogram per feature across
|
||||
all episodes rather than aggregating per-episode quantile summaries. The
|
||||
resulting quantiles are histogram approximations, subject to discretization and
|
||||
range-rebinning error; image/video frames are sampled by default.
|
||||
|
||||
Usage:
|
||||
|
||||
```bash
|
||||
@@ -39,7 +34,9 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import concurrent.futures
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
@@ -52,10 +49,11 @@ from lerobot.datasets import (
|
||||
CODEBASE_VERSION,
|
||||
DEFAULT_QUANTILES,
|
||||
LeRobotDataset,
|
||||
aggregate_stats,
|
||||
get_feature_stats,
|
||||
write_stats,
|
||||
)
|
||||
from lerobot.datasets.compute_stats import RunningQuantileStats, sample_indices
|
||||
from lerobot.datasets.compute_stats import sample_indices
|
||||
from lerobot.utils.utils import init_logging
|
||||
|
||||
|
||||
@@ -81,25 +79,20 @@ def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[s
|
||||
return False
|
||||
|
||||
|
||||
def collect_episode_arrays(
|
||||
dataset: LeRobotDataset,
|
||||
episode_idx: int,
|
||||
use_sampling: bool = True,
|
||||
skip_images: bool = False,
|
||||
) -> dict[str, tuple[np.ndarray, int]]:
|
||||
"""Collect one episode's frames per feature, flattened to (num_samples, dim).
|
||||
def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampling: bool = True) -> dict:
|
||||
"""Process a single episode and return its statistics.
|
||||
|
||||
Args:
|
||||
dataset: The LeRobot dataset
|
||||
episode_idx: Index of the episode to read
|
||||
use_sampling: If True, sub-sample image/video frames to bound memory.
|
||||
If False, use every frame (higher memory).
|
||||
skip_images: If True, skip image/video features entirely.
|
||||
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:
|
||||
Mapping of feature name to that episode's values and the number of frames
|
||||
they came from (which differs from the row count for image features).
|
||||
Dictionary containing episode statistics
|
||||
"""
|
||||
logging.info(f"Computing stats for episode {episode_idx}")
|
||||
|
||||
start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
|
||||
end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]
|
||||
|
||||
@@ -109,9 +102,7 @@ def collect_episode_arrays(
|
||||
# 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", "language")
|
||||
k for k in dataset.features if dataset.features[k]["dtype"] not in ("image", "video", "string")
|
||||
]
|
||||
|
||||
collected_data: dict[str, list] = {}
|
||||
@@ -123,7 +114,7 @@ def collect_episode_arrays(
|
||||
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 and not skip_images:
|
||||
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]
|
||||
@@ -131,82 +122,87 @@ def collect_episode_arrays(
|
||||
if key in item:
|
||||
collected_data.setdefault(key, []).append(item[key])
|
||||
|
||||
episode_arrays: dict[str, tuple[np.ndarray, int]] = {}
|
||||
ep_stats = {}
|
||||
for key, data_list in collected_data.items():
|
||||
if dataset.features[key]["dtype"] == "string":
|
||||
continue
|
||||
|
||||
data = torch.stack(data_list).cpu().numpy()
|
||||
if dataset.features[key]["dtype"] in ["image", "video"]:
|
||||
if data.dtype == np.uint8:
|
||||
data = data.astype(np.float32) / 255.0
|
||||
# (N, C, H, W) -> (N * H * W, C) so quantiles are computed per channel.
|
||||
channels = data.shape[1]
|
||||
values = data.transpose(0, 2, 3, 1).reshape(-1, channels)
|
||||
|
||||
axes_to_reduce = (0, 2, 3)
|
||||
keepdims = True
|
||||
else:
|
||||
values = data.reshape(-1, data.shape[-1]) if data.ndim > 1 else data.reshape(-1, 1)
|
||||
episode_arrays[key] = (values, len(data_list))
|
||||
axes_to_reduce = 0
|
||||
keepdims = data.ndim == 1
|
||||
|
||||
return episode_arrays
|
||||
ep_stats[key] = get_feature_stats(
|
||||
data, axis=axes_to_reduce, keepdims=keepdims, quantile_list=DEFAULT_QUANTILES
|
||||
)
|
||||
|
||||
if dataset.features[key]["dtype"] in ["image", "video"]:
|
||||
ep_stats[key] = {
|
||||
k: v if k == "count" else np.squeeze(v, axis=0) for k, v in ep_stats[key].items()
|
||||
}
|
||||
|
||||
return ep_stats
|
||||
|
||||
|
||||
def compute_quantile_stats_for_dataset(
|
||||
dataset: LeRobotDataset,
|
||||
use_sampling: bool = True,
|
||||
skip_images: bool = False,
|
||||
) -> dict[str, dict]:
|
||||
"""Compute whole-dataset statistics with one running histogram per feature.
|
||||
def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]:
|
||||
"""Compute quantile statistics for all episodes in the dataset.
|
||||
|
||||
Args:
|
||||
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 (higher memory).
|
||||
skip_images: If True, skip image/video features and leave their stats untouched.
|
||||
memory. If False, use every frame (exact, higher memory).
|
||||
|
||||
Returns:
|
||||
Dictionary containing statistics with histogram-based global quantile estimates
|
||||
Dictionary containing aggregated statistics with quantiles
|
||||
|
||||
Note:
|
||||
Episodes are accumulated sequentially because the running accumulators are
|
||||
shared across all of them.
|
||||
Video decoding operations are not thread-safe, so we process episodes sequentially
|
||||
when video keys are present. For datasets without videos, we use parallel processing
|
||||
with ThreadPoolExecutor for better performance.
|
||||
"""
|
||||
logging.info(f"Computing quantile statistics for dataset with {dataset.num_episodes} episodes")
|
||||
|
||||
running_stats: dict[str, RunningQuantileStats] = {}
|
||||
frame_counts: dict[str, int] = {}
|
||||
row_counts: dict[str, int] = {}
|
||||
# Kept only while a feature has a single row, so it can still be finalized.
|
||||
single_row_arrays: dict[str, np.ndarray] = {}
|
||||
episode_stats_list = []
|
||||
has_videos = len(dataset.meta.video_keys) > 0
|
||||
|
||||
for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
|
||||
episode_arrays = collect_episode_arrays(
|
||||
dataset, episode_idx, use_sampling=use_sampling, skip_images=skip_images
|
||||
)
|
||||
for key, (array, num_frames) in episode_arrays.items():
|
||||
running_stats.setdefault(key, RunningQuantileStats()).update(array)
|
||||
frame_counts[key] = frame_counts.get(key, 0) + num_frames
|
||||
row_counts[key] = row_counts.get(key, 0) + len(array)
|
||||
if row_counts[key] < 2:
|
||||
single_row_arrays[key] = array
|
||||
else:
|
||||
single_row_arrays.pop(key, None)
|
||||
if has_videos:
|
||||
logging.info("Dataset contains video keys - using sequential processing for thread safety")
|
||||
for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
|
||||
ep_stats = process_single_episode(dataset, episode_idx, use_sampling)
|
||||
episode_stats_list.append(ep_stats)
|
||||
else:
|
||||
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)))
|
||||
|
||||
if not running_stats:
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
future_to_episode = {
|
||||
executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx
|
||||
for episode_idx in range(dataset.num_episodes)
|
||||
}
|
||||
|
||||
episode_results = {}
|
||||
with tqdm(total=dataset.num_episodes, desc="Processing episodes") as pbar:
|
||||
for future in concurrent.futures.as_completed(future_to_episode):
|
||||
episode_idx = future_to_episode[future]
|
||||
ep_stats = future.result()
|
||||
episode_results[episode_idx] = ep_stats
|
||||
pbar.update(1)
|
||||
|
||||
for episode_idx in range(dataset.num_episodes):
|
||||
if episode_idx in episode_results:
|
||||
episode_stats_list.append(episode_results[episode_idx])
|
||||
|
||||
if not episode_stats_list:
|
||||
raise ValueError("No episode data found for computing statistics")
|
||||
|
||||
aggregated_stats: dict[str, dict] = {}
|
||||
for key, accumulator in running_stats.items():
|
||||
if row_counts[key] < 2:
|
||||
# Histograms need at least two samples; mirror get_feature_stats' basic-stats path.
|
||||
stats = get_feature_stats(single_row_arrays[key], axis=0, keepdims=False)
|
||||
else:
|
||||
stats = accumulator.get_statistics()
|
||||
if dataset.features[key]["dtype"] in ["image", "video"]:
|
||||
# Image stats are stored as (C, 1, 1) to broadcast over height and width.
|
||||
stats = {k: v if k == "count" else v[:, np.newaxis, np.newaxis] for k, v in stats.items()}
|
||||
# `get_feature_stats` counts frames, not the per-channel rows the accumulator sees.
|
||||
stats["count"] = np.array([frame_counts[key]])
|
||||
aggregated_stats[key] = stats
|
||||
|
||||
logging.info(f"Computed global histogram statistics for {len(aggregated_stats)} features")
|
||||
return aggregated_stats
|
||||
logging.info(f"Aggregating statistics from {len(episode_stats_list)} episodes")
|
||||
return aggregate_stats(episode_stats_list)
|
||||
|
||||
|
||||
def augment_dataset_with_quantile_stats(
|
||||
@@ -214,7 +210,6 @@ def augment_dataset_with_quantile_stats(
|
||||
root: str | Path | None = None,
|
||||
overwrite: bool = False,
|
||||
use_sampling: bool = True,
|
||||
skip_images: bool = False,
|
||||
) -> None:
|
||||
"""Augment a dataset with quantile statistics if they are missing.
|
||||
|
||||
@@ -223,8 +218,7 @@ def augment_dataset_with_quantile_stats(
|
||||
root: Local root directory for the dataset
|
||||
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 (higher memory).
|
||||
skip_images: If True, skip image/video features and keep their existing stats
|
||||
memory. If False, use every frame (exact, higher memory).
|
||||
"""
|
||||
logging.info(f"Loading dataset: {repo_id}")
|
||||
dataset = LeRobotDataset(
|
||||
@@ -238,13 +232,7 @@ def augment_dataset_with_quantile_stats(
|
||||
|
||||
logging.info("Dataset does not contain quantile statistics. Computing them now...")
|
||||
|
||||
new_stats = compute_quantile_stats_for_dataset(
|
||||
dataset, use_sampling=use_sampling, skip_images=skip_images
|
||||
)
|
||||
|
||||
if skip_images and dataset.meta.stats:
|
||||
for key, feature_stats in dataset.meta.stats.items():
|
||||
new_stats.setdefault(key, feature_stats)
|
||||
new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling)
|
||||
|
||||
logging.info("Updating dataset metadata with new quantile statistics")
|
||||
dataset.meta.stats = new_stats
|
||||
@@ -288,15 +276,10 @@ def main():
|
||||
"--no-sampling",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Compute stats over every frame (higher memory). By default, "
|
||||
"Compute stats over every frame (exact, higher memory). By default, "
|
||||
"image/video frames are sub-sampled per episode to bound memory."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-images",
|
||||
action="store_true",
|
||||
help="Skip image/video features and preserve their existing stats",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
root = Path(args.root) if args.root else None
|
||||
@@ -308,7 +291,6 @@ def main():
|
||||
root=root,
|
||||
overwrite=args.overwrite,
|
||||
use_sampling=not args.no_sampling,
|
||||
skip_images=args.skip_images,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -76,7 +76,7 @@ from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_proces
|
||||
from lerobot.policies.factory import ProcessorConfigKwargs
|
||||
from lerobot.rewards import make_reward_pre_post_processors
|
||||
from lerobot.utils.collate import lerobot_collate_fn
|
||||
from lerobot.utils.constants import PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR
|
||||
from lerobot.utils.constants import TRAINING_STATE_DIR
|
||||
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
|
||||
from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
|
||||
from lerobot.utils.random_utils import set_seed
|
||||
@@ -95,20 +95,6 @@ else:
|
||||
|
||||
from .lerobot_eval import eval_policy_all
|
||||
|
||||
EMA_STATE_FILENAME = "ema_state.pt"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _ema_weights(ema: Any, policy: PreTrainedPolicy) -> Iterator[None]:
|
||||
"""Temporarily swap the EMA shadow weights into `policy`, restoring the live ones on exit."""
|
||||
params = list(policy.parameters())
|
||||
ema.store(params)
|
||||
ema.copy_to(params)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
ema.restore(params)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _make_eval_envs(cfg: TrainPipelineConfig) -> Iterator[dict[str, dict[int, Any]]]:
|
||||
@@ -606,65 +592,6 @@ def train(cfg: TrainPipelineConfig):
|
||||
dl_iter = cycle(dataloader)
|
||||
policy.train()
|
||||
|
||||
# EMA shadow of the policy weights (Chi et al. 2023, Diffusion Policy, section V.D). The shadow
|
||||
# lives on the main process only, which is safe under DDP where every rank holds identical
|
||||
# weights after each gradient sync. diffusers is imported lazily so the base training path does
|
||||
# not depend on it.
|
||||
ema = None
|
||||
if cfg.ema.enable:
|
||||
if parallel_dims.is_sharded:
|
||||
raise NotImplementedError(
|
||||
"--ema.enable=true is not supported with sharded training (FSDP2/HSDP/CP): the "
|
||||
"parameters are sharded across ranks. Use a replicated (DDP) or single-GPU run."
|
||||
)
|
||||
if cfg.peft is not None:
|
||||
raise NotImplementedError("--ema.enable=true is not supported together with PEFT adapters.")
|
||||
require_package("diffusers", extra="diffusion")
|
||||
if is_main_process():
|
||||
from diffusers.training_utils import EMAModel # noqa: PLC0415
|
||||
|
||||
# A constant --ema.decay is expressed through the schedule clamp: with
|
||||
# min_decay == max_decay, the warmup curve is pinned to that value at every step.
|
||||
min_decay = cfg.ema.min_decay if cfg.ema.decay is None else cfg.ema.decay
|
||||
max_decay = cfg.ema.max_decay if cfg.ema.decay is None else cfg.ema.decay
|
||||
ema = EMAModel(
|
||||
accelerator.unwrap_model(policy).parameters(),
|
||||
decay=max_decay,
|
||||
min_decay=min_decay,
|
||||
update_after_step=cfg.ema.update_after_step,
|
||||
use_ema_warmup=True,
|
||||
inv_gamma=cfg.ema.inv_gamma,
|
||||
power=cfg.ema.power,
|
||||
)
|
||||
ema.to(device)
|
||||
if cfg.ema.decay is not None:
|
||||
logging.info(
|
||||
"EMA enabled: decay=%g (constant), update_after_step=%d, use_for_eval=%s",
|
||||
cfg.ema.decay,
|
||||
cfg.ema.update_after_step,
|
||||
cfg.ema.use_for_eval,
|
||||
)
|
||||
else:
|
||||
logging.info(
|
||||
"EMA enabled: max_decay=%g, inv_gamma=%g, power=%g, update_after_step=%d, use_for_eval=%s",
|
||||
cfg.ema.max_decay,
|
||||
cfg.ema.inv_gamma,
|
||||
cfg.ema.power,
|
||||
cfg.ema.update_after_step,
|
||||
cfg.ema.use_for_eval,
|
||||
)
|
||||
if cfg.checkpoint_path is not None:
|
||||
ema_path = cfg.checkpoint_path / TRAINING_STATE_DIR / EMA_STATE_FILENAME
|
||||
if ema_path.exists():
|
||||
ema.load_state_dict(torch.load(ema_path, map_location=device, weights_only=True))
|
||||
logging.info("Resumed EMA shadow from %s", ema_path)
|
||||
else:
|
||||
logging.warning(
|
||||
"Resuming with --ema.enable=true but %s is missing; "
|
||||
"restarting the shadow from the current weights.",
|
||||
ema_path,
|
||||
)
|
||||
|
||||
train_metrics = {
|
||||
# Per-rank loss reflects only one shard of the global batch; mean recovers the loss the
|
||||
# data-parallel group is actually optimizing. grad_norm and lr are already identical on
|
||||
@@ -675,10 +602,9 @@ def train(cfg: TrainPipelineConfig):
|
||||
"lr": AverageMeter("lr", ":0.1e"),
|
||||
# Report the slowest rank for bottleneck-style timings so multi-GPU runs surface the
|
||||
# true straggler instead of rank 0's view.
|
||||
"dataloading_s": AverageMeter("data_s", ":.3f", reduction="max"),
|
||||
"preprocessing_s": AverageMeter("prep_s", ":.3f", reduction="max"),
|
||||
"update_s": AverageMeter("updt_s", ":.3f", reduction="max"),
|
||||
"step_s": AverageMeter("step_s", ":.3f", reduction="max"),
|
||||
"dataloading_s": AverageMeter("data_s", ":.3f", reduction="max"),
|
||||
# Derived from the post-reduce max step time; set once per log window on the main rank.
|
||||
"samples_per_s": AverageMeter("smp/s", ":.0f"),
|
||||
}
|
||||
if torch.cuda.is_available():
|
||||
@@ -708,15 +634,13 @@ def train(cfg: TrainPipelineConfig):
|
||||
)
|
||||
|
||||
for _ in range(step, cfg.steps):
|
||||
step_start = time.perf_counter()
|
||||
start_time = time.perf_counter()
|
||||
batch = next(dl_iter)
|
||||
preprocessing_start = time.perf_counter()
|
||||
train_tracker.dataloading_s = preprocessing_start - step_start
|
||||
for cam_key in dataset.meta.camera_keys:
|
||||
if cam_key in batch and batch[cam_key].dtype == torch.uint8:
|
||||
batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0
|
||||
batch = preprocessor(batch)
|
||||
train_tracker.preprocessing_s = time.perf_counter() - preprocessing_start
|
||||
train_tracker.dataloading_s = time.perf_counter() - start_time
|
||||
|
||||
train_tracker, _ = update_policy(
|
||||
train_tracker,
|
||||
@@ -728,13 +652,6 @@ def train(cfg: TrainPipelineConfig):
|
||||
lr_scheduler=lr_scheduler,
|
||||
sample_weighter=sample_weighter,
|
||||
)
|
||||
train_tracker.step_s = time.perf_counter() - step_start
|
||||
|
||||
# Pull one optimizer step of the live weights into the EMA shadow (main process only).
|
||||
# The shadow tracks optimizer updates, not micro-batches: gate on the sync step under
|
||||
# gradient accumulation.
|
||||
if ema is not None and accelerator.sync_gradients:
|
||||
ema.step(accelerator.unwrap_model(policy).parameters())
|
||||
|
||||
# Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
|
||||
# increment `step` here.
|
||||
@@ -751,8 +668,11 @@ def train(cfg: TrainPipelineConfig):
|
||||
# Collective reduce must run on every rank, before the main-process gate below.
|
||||
train_tracker.reduce_across_ranks()
|
||||
if is_main_process():
|
||||
if train_tracker.step_s.avg > 0:
|
||||
train_tracker.samples_per_s = samples_per_step / train_tracker.step_s.avg
|
||||
# Cluster-wide throughput, derived from the already-reduced (max) step time so it
|
||||
# reflects the slowest rank — which is what actually gates the next iteration.
|
||||
step_time = train_tracker.update_s.avg + train_tracker.dataloading_s.avg
|
||||
if step_time > 0:
|
||||
train_tracker.samples_per_s = samples_per_step / step_time
|
||||
logging.info(train_tracker)
|
||||
if wandb_logger:
|
||||
# Policy sub-losses (latent_loss, action_loss, ...) are aggregated into the
|
||||
@@ -763,9 +683,6 @@ def train(cfg: TrainPipelineConfig):
|
||||
if sample_weighter is not None:
|
||||
weighter_stats = sample_weighter.get_stats()
|
||||
wandb_log_dict.update({f"sample_weighting/{k}": v for k, v in weighter_stats.items()})
|
||||
if ema is not None and ema.cur_decay_value is not None:
|
||||
wandb_log_dict["ema/decay"] = ema.cur_decay_value
|
||||
wandb_log_dict["ema/step"] = ema.optimization_step
|
||||
wandb_logger.log_dict(wandb_log_dict, step)
|
||||
train_tracker.reset_averages()
|
||||
|
||||
@@ -810,17 +727,6 @@ def train(cfg: TrainPipelineConfig):
|
||||
accelerator=accelerator,
|
||||
)
|
||||
if is_main_process():
|
||||
if ema is not None:
|
||||
# Save the shadow for exact resume, plus a directly loadable copy of the EMA
|
||||
# weights (lerobot-eval --policy.path=<checkpoint>/pretrained_model_ema).
|
||||
torch.save(ema.state_dict(), checkpoint_dir / TRAINING_STATE_DIR / EMA_STATE_FILENAME)
|
||||
unwrapped_policy = accelerator.unwrap_model(policy)
|
||||
ema_dir = checkpoint_dir / f"{PRETRAINED_MODEL_DIR}_ema"
|
||||
with _ema_weights(ema, unwrapped_policy):
|
||||
unwrapped_policy.save_pretrained(ema_dir)
|
||||
cfg.save_pretrained(ema_dir)
|
||||
preprocessor.save_pretrained(ema_dir)
|
||||
postprocessor.save_pretrained(ema_dir)
|
||||
update_last_checkpoint(checkpoint_dir)
|
||||
if cfg.save_checkpoint_to_hub:
|
||||
push_checkpoint_to_hub(
|
||||
@@ -836,18 +742,10 @@ def train(cfg: TrainPipelineConfig):
|
||||
if is_main_process():
|
||||
step_id = get_step_identifier(step, cfg.steps)
|
||||
logging.info(f"Eval policy at step {step}")
|
||||
eval_policy_model = accelerator.unwrap_model(policy)
|
||||
# Evaluate the EMA weights when enabled: the swap happens only on the main
|
||||
# process (the other ranks wait at the barrier below) and is exactly undone
|
||||
# afterwards, so the live weights stay in sync across ranks.
|
||||
use_ema_for_eval = ema is not None and cfg.ema.use_for_eval
|
||||
if use_ema_for_eval:
|
||||
logging.info("Evaluating the EMA weights")
|
||||
weights_cm = _ema_weights(ema, eval_policy_model) if use_ema_for_eval else nullcontext()
|
||||
with weights_cm, _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
|
||||
with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
|
||||
eval_info = eval_policy_all(
|
||||
envs=eval_env, # dict[suite][task_id] -> vec_env
|
||||
policy=eval_policy_model,
|
||||
policy=accelerator.unwrap_model(policy),
|
||||
env_preprocessor=env_preprocessor,
|
||||
env_postprocessor=env_postprocessor,
|
||||
preprocessor=preprocessor,
|
||||
@@ -906,25 +804,6 @@ def train(cfg: TrainPipelineConfig):
|
||||
peft_model=unwrapped if peft_model is not None else None,
|
||||
)
|
||||
|
||||
# The push above ships the live weights; when EMA is on, the weights that were
|
||||
# evaluated are the shadow, so push those too under a sibling `<repo_id>-ema` repo.
|
||||
# The shadow lives on the main process only, so this is rank-0-only by construction.
|
||||
# Non-fatal: the live model is already up if this fails.
|
||||
if ema is not None:
|
||||
ema_repo_id = f"{active_cfg.repo_id}-ema"
|
||||
orig_repo_id = unwrapped.config.repo_id
|
||||
try:
|
||||
unwrapped.config.repo_id = ema_repo_id
|
||||
with _ema_weights(ema, unwrapped):
|
||||
unwrapped.push_model_to_hub(cfg, dataset_meta=dataset.meta)
|
||||
preprocessor.push_to_hub(ema_repo_id)
|
||||
postprocessor.push_to_hub(ema_repo_id)
|
||||
logging.info("Pushed EMA weights to %s", ema_repo_id)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logging.warning("Failed to push EMA weights to %s: %s", ema_repo_id, exc)
|
||||
finally:
|
||||
unwrapped.config.repo_id = orig_repo_id
|
||||
|
||||
# Properly clean up the distributed process group
|
||||
accelerator.wait_for_everyone()
|
||||
accelerator.end_training()
|
||||
|
||||
@@ -35,8 +35,8 @@ class KeyboardEndEffectorTeleopConfig(KeyboardTeleopConfig):
|
||||
|
||||
Used for controlling robot end-effectors with keyboard inputs.
|
||||
|
||||
Attributes:
|
||||
use_gripper: Whether to include gripper control in actions
|
||||
**Attributes**:
|
||||
- **use_gripper** (`bool`) -- Whether to include gripper control in actions
|
||||
"""
|
||||
|
||||
use_gripper: bool = True
|
||||
@@ -49,14 +49,14 @@ class KeyboardRoverTeleopConfig(TeleoperatorConfig):
|
||||
|
||||
Used for controlling mobile robots like EarthRover Mini Plus with WASD controls.
|
||||
|
||||
Attributes:
|
||||
linear_speed: Default linear velocity magnitude (-1 to 1 range for SDK robots)
|
||||
angular_speed: Default angular velocity magnitude (-1 to 1 range for SDK robots)
|
||||
speed_increment: Amount to increase/decrease speed with +/- keys
|
||||
turn_assist_ratio: Forward motion multiplier when turning with A/D keys (0.0-1.0)
|
||||
angular_speed_ratio: Ratio of angular to linear speed for synchronized adjustments
|
||||
min_linear_speed: Minimum linear speed when decreasing (prevents zero speed)
|
||||
min_angular_speed: Minimum angular speed when decreasing (prevents zero speed)
|
||||
**Attributes**:
|
||||
- **linear_speed** (`float`) -- Default linear velocity magnitude (-1 to 1 range for SDK robots)
|
||||
- **angular_speed** (`float`) -- Default angular velocity magnitude (-1 to 1 range for SDK robots)
|
||||
- **speed_increment** (`float`) -- Amount to increase/decrease speed with +/- keys
|
||||
- **turn_assist_ratio** (`float`) -- Forward motion multiplier when turning with A/D keys (0.0-1.0)
|
||||
- **angular_speed_ratio** (`float`) -- Ratio of angular to linear speed for synchronized adjustments
|
||||
- **min_linear_speed** (`float`) -- Minimum linear speed when decreasing (prevents zero speed)
|
||||
- **min_angular_speed** (`float`) -- Minimum angular speed when decreasing (prevents zero speed)
|
||||
"""
|
||||
|
||||
linear_speed: float = 1.0
|
||||
|
||||
@@ -312,10 +312,10 @@ class KeyboardRoverTeleop(KeyboardTeleop):
|
||||
System:
|
||||
- ESC: Disconnect teleoperator
|
||||
|
||||
Attributes:
|
||||
config: Teleoperator configuration
|
||||
current_linear_speed: Current linear velocity magnitude
|
||||
current_angular_speed: Current angular velocity magnitude
|
||||
**Attributes**:
|
||||
- **config** -- Teleoperator configuration
|
||||
- **current_linear_speed** -- Current linear velocity magnitude
|
||||
- **current_angular_speed** -- Current angular velocity magnitude
|
||||
|
||||
Example:
|
||||
```python
|
||||
|
||||
@@ -35,9 +35,9 @@ class MapPhoneActionToRobotAction(RobotActionProcessorStep):
|
||||
necessary axis inversions and swaps. It also interprets platform-specific
|
||||
button presses to generate a gripper command.
|
||||
|
||||
Attributes:
|
||||
platform: The operating system of the phone (iOS or Android), used
|
||||
to determine the correct button mappings for the gripper.
|
||||
**Attributes**:
|
||||
- **platform** (`PhoneOS`) -- The operating system of the phone (iOS or Android), used to determine
|
||||
the correct button mappings for the gripper.
|
||||
"""
|
||||
|
||||
# TODO(Steven): Gripper vel could be output of phone_teleop directly
|
||||
|
||||
@@ -27,15 +27,22 @@ from .config import TeleoperatorConfig
|
||||
|
||||
|
||||
class Teleoperator(abc.ABC):
|
||||
"""
|
||||
The base abstract class for all LeRobot-compatible teleoperation devices.
|
||||
"""The base abstract class for all LeRobot-compatible teleoperation devices.
|
||||
|
||||
This class provides a standardized interface for interacting with physical teleoperators.
|
||||
Subclasses must implement all abstract methods and properties to be usable.
|
||||
This class provides a standardized interface for interacting with physical teleoperators. Subclasses
|
||||
must implement all abstract methods and properties to be usable.
|
||||
|
||||
Attributes:
|
||||
config_class (RobotConfig): The expected configuration class for this teleoperator.
|
||||
name (str): The unique name used to identify this teleoperator type.
|
||||
Used as a context manager, a teleoperator connects on entry and disconnects on exit, even on error:
|
||||
|
||||
```python
|
||||
>>> with SO101Leader(config) as teleop: # doctest: +SKIP
|
||||
... action = teleop.get_action()
|
||||
```
|
||||
|
||||
**Attributes**:
|
||||
- **config_class** (`type[TeleoperatorConfig]`) -- The expected configuration class for this
|
||||
teleoperator.
|
||||
- **name** (`str`) -- The unique name used to identify this teleoperator type.
|
||||
"""
|
||||
|
||||
# Set these in ALL subclasses
|
||||
@@ -59,25 +66,16 @@ class Teleoperator(abc.ABC):
|
||||
return f"{self.id} {self.__class__.__name__}"
|
||||
|
||||
def __enter__(self):
|
||||
"""
|
||||
Context manager entry.
|
||||
Automatically connects to the camera.
|
||||
"""
|
||||
"""Context manager entry. Automatically connects to the teleoperator."""
|
||||
self.connect()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
||||
"""
|
||||
Context manager exit.
|
||||
Automatically disconnects, ensuring resources are released even on error.
|
||||
"""
|
||||
"""Context manager exit. Disconnects, ensuring resources are released even on error."""
|
||||
self.disconnect()
|
||||
|
||||
def __del__(self) -> None:
|
||||
"""
|
||||
Destructor safety net.
|
||||
Attempts to disconnect if the object is garbage collected without cleanup.
|
||||
"""
|
||||
"""Destructor safety net. Disconnects if the object is garbage collected without cleanup."""
|
||||
try:
|
||||
if self.is_connected:
|
||||
self.disconnect()
|
||||
@@ -87,82 +85,98 @@ class Teleoperator(abc.ABC):
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def action_features(self) -> dict:
|
||||
"""
|
||||
A dictionary describing the structure and types of the actions produced by the teleoperator. Its
|
||||
structure (keys) should match the structure of what is returned by :pymeth:`get_action`. Values for
|
||||
the dict should be the type of the value if it's a simple value, e.g. `float` for single
|
||||
proprioceptive value (a joint's goal position/velocity)
|
||||
"""A dictionary describing the structure and types of the actions produced by the teleoperator.
|
||||
|
||||
Note: this property should be able to be called regardless of whether the robot is connected or not.
|
||||
Its keys should match the structure of what is returned by [`~teleoperators.Teleoperator.get_action`].
|
||||
Values should be the type of the value if it's a simple value, e.g. `float` for a single
|
||||
proprioceptive value (a joint's goal position or velocity).
|
||||
|
||||
> [!NOTE]
|
||||
> This property must be callable regardless of whether the teleoperator is connected.
|
||||
|
||||
Returns:
|
||||
`dict`: Action names mapped to their type or shape.
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def feedback_features(self) -> dict:
|
||||
"""
|
||||
A dictionary describing the structure and types of the feedback actions expected by the robot. Its
|
||||
structure (keys) should match the structure of what is passed to :pymeth:`send_feedback`. Values for
|
||||
the dict should be the type of the value if it's a simple value, e.g. `float` for single
|
||||
proprioceptive value (a joint's goal position/velocity)
|
||||
"""A dictionary describing the structure and types of the feedback actions the teleoperator accepts.
|
||||
|
||||
Note: this property should be able to be called regardless of whether the robot is connected or not.
|
||||
Its keys should match the structure of what is passed to
|
||||
[`~teleoperators.Teleoperator.send_feedback`]. Values should be the type of the value if it's a
|
||||
simple value, e.g. `float` for a single proprioceptive value (a joint's goal position or velocity).
|
||||
|
||||
> [!NOTE]
|
||||
> This property must be callable regardless of whether the teleoperator is connected.
|
||||
|
||||
Returns:
|
||||
`dict`: Feedback names mapped to their type or shape.
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def is_connected(self) -> bool:
|
||||
"""
|
||||
Whether the teleoperator is currently connected or not. If `False`, calling :pymeth:`get_action`
|
||||
or :pymeth:`send_feedback` should raise an error.
|
||||
"""Whether the teleoperator is currently connected.
|
||||
|
||||
If `False`, calling [`~teleoperators.Teleoperator.get_action`] or
|
||||
[`~teleoperators.Teleoperator.send_feedback`] should raise an error.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` if communication with the teleoperator is established.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""
|
||||
Establish communication with the teleoperator.
|
||||
"""Establish communication with the teleoperator.
|
||||
|
||||
Args:
|
||||
calibrate (bool): If True, automatically calibrate the teleoperator after connecting if it's not
|
||||
calibrated or needs calibration (this is hardware-dependant).
|
||||
calibrate (`bool`, *optional*, defaults to `True`):
|
||||
Whether to automatically calibrate the teleoperator after connecting, if it is not
|
||||
calibrated or needs recalibration. Whether calibration is needed is hardware-dependent.
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def is_calibrated(self) -> bool:
|
||||
"""Whether the teleoperator is currently calibrated or not. Should be always `True` if not applicable"""
|
||||
"""Whether the teleoperator is currently calibrated.
|
||||
|
||||
Returns:
|
||||
`bool`: `True` if the teleoperator is calibrated. Always `True` for teleoperators where
|
||||
calibration does not apply.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def calibrate(self) -> None:
|
||||
"""
|
||||
Calibrate the teleoperator if applicable. If not, this should be a no-op.
|
||||
"""Calibrate the teleoperator if applicable. If not, this should be a no-op.
|
||||
|
||||
This method should collect any necessary data (e.g., motor offsets) and update the
|
||||
:pyattr:`calibration` dictionary accordingly.
|
||||
This method should collect any necessary data (e.g. motor offsets) and update the `calibration`
|
||||
dictionary accordingly.
|
||||
"""
|
||||
pass
|
||||
|
||||
def _load_calibration(self, fpath: Path | None = None) -> None:
|
||||
"""
|
||||
Helper to load calibration data from the specified file.
|
||||
"""Helper to load calibration data from the specified file.
|
||||
|
||||
Args:
|
||||
fpath (Path | None): Optional path to the calibration file. Defaults to `self.calibration_fpath`.
|
||||
fpath (`Path`, *optional*):
|
||||
Path to the calibration file. Defaults to `self.calibration_fpath`.
|
||||
"""
|
||||
fpath = self.calibration_fpath if fpath is None else fpath
|
||||
with open(fpath) as f, draccus.config_type("json"):
|
||||
self.calibration = draccus.load(dict[str, MotorCalibration], f)
|
||||
|
||||
def _save_calibration(self, fpath: Path | None = None) -> None:
|
||||
"""
|
||||
Helper to save calibration data to the specified file.
|
||||
"""Helper to save calibration data to the specified file.
|
||||
|
||||
Args:
|
||||
fpath (Path | None): Optional path to save the calibration file. Defaults to `self.calibration_fpath`.
|
||||
fpath (`Path`, *optional*):
|
||||
Path to save the calibration file to. Defaults to `self.calibration_fpath`.
|
||||
"""
|
||||
fpath = self.calibration_fpath if fpath is None else fpath
|
||||
with open(fpath, "w") as f, draccus.config_type("json"):
|
||||
@@ -170,35 +184,36 @@ class Teleoperator(abc.ABC):
|
||||
|
||||
@abc.abstractmethod
|
||||
def configure(self) -> None:
|
||||
"""
|
||||
Apply any one-time or runtime configuration to the teleoperator.
|
||||
"""Apply any one-time or runtime configuration to the teleoperator.
|
||||
|
||||
This may include setting motor parameters, control modes, or initial state.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def get_action(self) -> RobotAction:
|
||||
"""
|
||||
Retrieve the current action from the teleoperator.
|
||||
"""Retrieve the current action from the teleoperator.
|
||||
|
||||
Returns:
|
||||
RobotAction: A flat dictionary representing the teleoperator's current actions. Its
|
||||
structure should match :pymeth:`observation_features`.
|
||||
`dict[str, Any]`: A flat dictionary representing the teleoperator's current action. Its structure
|
||||
should match [`~teleoperators.Teleoperator.action_features`].
|
||||
|
||||
Raises:
|
||||
DeviceNotConnectedError: If [`~teleoperators.Teleoperator.connect`] has not been called.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def send_feedback(self, feedback: dict[str, Any]) -> None:
|
||||
"""
|
||||
Send a feedback action command to the teleoperator.
|
||||
"""Send a feedback command to the teleoperator, e.g. force feedback on a leader arm.
|
||||
|
||||
Args:
|
||||
feedback (dict[str, Any]): Dictionary representing the desired feedback. Its structure should match
|
||||
:pymeth:`feedback_features`.
|
||||
feedback (`dict[str, Any]`):
|
||||
The desired feedback. Its structure should match
|
||||
[`~teleoperators.Teleoperator.feedback_features`].
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: The action actually sent to the motors potentially clipped or modified, e.g. by
|
||||
safety limits on velocity.
|
||||
Raises:
|
||||
DeviceNotConnectedError: If [`~teleoperators.Teleoperator.connect`] has not been called.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@@ -82,13 +82,13 @@ class SampleWeightingConfig:
|
||||
The `type` field determines which implementation to use, and `extra_params`
|
||||
contains additional type-specific parameters.
|
||||
|
||||
Attributes:
|
||||
type: Weighting strategy type ("rabc", "uniform", etc.)
|
||||
progress_path: Path to precomputed progress values (for RABC)
|
||||
head_mode: Which model head to use for progress ("sparse" or "dense")
|
||||
kappa: Hard threshold for high-quality samples (RABC-specific)
|
||||
epsilon: Small constant for numerical stability
|
||||
extra_params: Additional type-specific parameters passed to the weighter
|
||||
**Attributes**:
|
||||
- **type** (`str`) -- Weighting strategy type ("rabc", "uniform", etc.)
|
||||
- **progress_path** (`str | None`) -- Path to precomputed progress values (for RABC)
|
||||
- **head_mode** (`str`) -- Which model head to use for progress ("sparse" or "dense")
|
||||
- **kappa** (`float`) -- Hard threshold for high-quality samples (RABC-specific)
|
||||
- **epsilon** (`float`) -- Small constant for numerical stability
|
||||
- **extra_params** (`dict`) -- Additional type-specific parameters passed to the weighter
|
||||
"""
|
||||
|
||||
type: str = "rabc"
|
||||
|
||||
@@ -12,11 +12,8 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
|
||||
@@ -27,9 +24,7 @@ from lerobot.scripts.augment_dataset_quantile_stats import (
|
||||
|
||||
|
||||
def _numeric_keys(dataset):
|
||||
return [
|
||||
k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string", "language")
|
||||
]
|
||||
return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")]
|
||||
|
||||
|
||||
def _image_keys(dataset):
|
||||
@@ -107,112 +102,3 @@ def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory)
|
||||
)
|
||||
stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
|
||||
assert has_quantile_stats(stats)
|
||||
|
||||
|
||||
class FakeHFDataset:
|
||||
"""Minimal stand-in exposing the column slicing used by the augment script."""
|
||||
|
||||
def __init__(self, columns: dict[str, list]):
|
||||
self._columns = columns
|
||||
|
||||
def select_columns(self, keys):
|
||||
return FakeHFDataset({key: self._columns[key] for key in keys})
|
||||
|
||||
def __getitem__(self, index):
|
||||
return {key: values[index] for key, values in self._columns.items()}
|
||||
|
||||
|
||||
def test_compute_quantile_stats_skips_language_features():
|
||||
class FakeDataset:
|
||||
num_episodes = 1
|
||||
features = {
|
||||
"action": {"dtype": "float32"},
|
||||
"observation.language": {"dtype": "language"},
|
||||
}
|
||||
meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}])
|
||||
hf_dataset = FakeHFDataset(
|
||||
{
|
||||
"action": [[0.0], [1.0]],
|
||||
"observation.language": [
|
||||
[{"role": "user", "content": "pick"}],
|
||||
[{"role": "assistant", "content": "done"}],
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
stats = compute_quantile_stats_for_dataset(FakeDataset())
|
||||
|
||||
assert set(stats) == {"action"}
|
||||
|
||||
|
||||
def test_compute_quantile_stats_skip_images_avoids_decoding():
|
||||
class FakeDataset:
|
||||
num_episodes = 1
|
||||
features = {
|
||||
"action": {"dtype": "float32"},
|
||||
"observation.images.cam": {"dtype": "video"},
|
||||
}
|
||||
meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}])
|
||||
hf_dataset = FakeHFDataset({"action": [[0.0], [1.0]]})
|
||||
|
||||
def __getitem__(self, index):
|
||||
raise AssertionError(f"video frame {index} was decoded despite skip_images=True")
|
||||
|
||||
stats = compute_quantile_stats_for_dataset(FakeDataset(), skip_images=True)
|
||||
|
||||
assert set(stats) == {"action"}
|
||||
|
||||
|
||||
def test_compute_quantile_stats_handles_single_frame():
|
||||
class FakeDataset:
|
||||
num_episodes = 1
|
||||
features = {"action": {"dtype": "float32"}}
|
||||
meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 1}])
|
||||
hf_dataset = FakeHFDataset({"action": [[5.0, 7.0]]})
|
||||
|
||||
stats = compute_quantile_stats_for_dataset(FakeDataset())
|
||||
|
||||
np.testing.assert_array_equal(stats["action"]["count"], np.array([1]))
|
||||
for key in ("min", "max", "mean", "q01", "q10", "q50", "q90", "q99"):
|
||||
np.testing.assert_allclose(stats["action"][key], np.array([5.0, 7.0]))
|
||||
|
||||
|
||||
def test_compute_quantile_stats_image_count_uses_frames():
|
||||
frames = [torch.zeros(3, 2, 2), torch.ones(3, 2, 2)]
|
||||
|
||||
class FakeDataset:
|
||||
num_episodes = 1
|
||||
features = {"observation.images.cam": {"dtype": "video"}}
|
||||
meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}])
|
||||
hf_dataset = FakeHFDataset({})
|
||||
|
||||
def __getitem__(self, index):
|
||||
return {"observation.images.cam": frames[index]}
|
||||
|
||||
stats = compute_quantile_stats_for_dataset(FakeDataset(), use_sampling=False)
|
||||
image_stats = stats["observation.images.cam"]
|
||||
|
||||
np.testing.assert_array_equal(image_stats["count"], np.array([2]))
|
||||
assert image_stats["mean"].shape == (3, 1, 1)
|
||||
np.testing.assert_allclose(image_stats["mean"], np.full((3, 1, 1), 0.5))
|
||||
|
||||
|
||||
def test_compute_quantile_stats_accumulates_across_episodes():
|
||||
values = [[float(value)] for value in range(100)] + [[float(value)] for value in range(1000, 1010)]
|
||||
|
||||
class FakeDataset:
|
||||
num_episodes = 2
|
||||
features = {"action": {"dtype": "float32"}}
|
||||
meta = SimpleNamespace(
|
||||
episodes=[
|
||||
{"dataset_from_index": 0, "dataset_to_index": 100},
|
||||
{"dataset_from_index": 100, "dataset_to_index": 110},
|
||||
]
|
||||
)
|
||||
hf_dataset = FakeHFDataset({"action": values})
|
||||
|
||||
stats = compute_quantile_stats_for_dataset(FakeDataset())
|
||||
|
||||
np.testing.assert_array_equal(stats["action"]["count"], np.array([110]))
|
||||
expected_q90 = np.percentile(np.asarray(values), 90, axis=0)
|
||||
np.testing.assert_allclose(stats["action"]["q90"], expected_q90, atol=0.1)
|
||||
|
||||
@@ -688,7 +688,7 @@ def test_compute_episode_stats_string_features_skipped():
|
||||
|
||||
|
||||
def test_aggregate_feature_stats_with_quantiles():
|
||||
"""Test aggregating feature stats that include quantiles uses conservative bounds."""
|
||||
"""Test aggregating feature stats that include quantiles."""
|
||||
stats_ft_list = [
|
||||
{
|
||||
"min": np.array([1.0]),
|
||||
@@ -697,9 +697,6 @@ def test_aggregate_feature_stats_with_quantiles():
|
||||
"std": np.array([2.0]),
|
||||
"count": np.array([100]),
|
||||
"q01": np.array([1.5]),
|
||||
"q10": np.array([2.0]),
|
||||
"q50": np.array([5.0]),
|
||||
"q90": np.array([9.0]),
|
||||
"q99": np.array([9.5]),
|
||||
},
|
||||
{
|
||||
@@ -709,21 +706,22 @@ def test_aggregate_feature_stats_with_quantiles():
|
||||
"std": np.array([2.5]),
|
||||
"count": np.array([150]),
|
||||
"q01": np.array([2.5]),
|
||||
"q10": np.array([3.0]),
|
||||
"q50": np.array([6.0]),
|
||||
"q90": np.array([11.0]),
|
||||
"q99": np.array([11.5]),
|
||||
},
|
||||
]
|
||||
|
||||
result = aggregate_feature_stats(stats_ft_list)
|
||||
|
||||
# Lower quantiles use min; upper quantiles use max, regardless of counts.
|
||||
np.testing.assert_allclose(result["q01"], np.array([1.5]), atol=1e-6)
|
||||
np.testing.assert_allclose(result["q10"], np.array([2.0]), atol=1e-6)
|
||||
np.testing.assert_allclose(result["q50"], np.array([5.0]), atol=1e-6)
|
||||
np.testing.assert_allclose(result["q90"], np.array([11.0]), atol=1e-6)
|
||||
np.testing.assert_allclose(result["q99"], np.array([11.5]), atol=1e-6)
|
||||
# Should preserve quantiles
|
||||
assert "q01" in result
|
||||
assert "q99" in result
|
||||
|
||||
# Verify quantile aggregation (weighted average)
|
||||
expected_q01 = (1.5 * 100 + 2.5 * 150) / 250 # ≈ 2.1
|
||||
expected_q99 = (9.5 * 100 + 11.5 * 150) / 250 # ≈ 10.7
|
||||
|
||||
np.testing.assert_allclose(result["q01"], np.array([expected_q01]), atol=1e-6)
|
||||
np.testing.assert_allclose(result["q99"], np.array([expected_q99]), atol=1e-6)
|
||||
|
||||
|
||||
def test_aggregate_stats_mixed_quantiles():
|
||||
@@ -880,60 +878,3 @@ def test_fixed_quantiles_always_computed():
|
||||
for q_key in expected_quantiles:
|
||||
assert q_key in episode_stats[key]
|
||||
assert episode_stats[key][q_key].shape == (features[key]["shape"][0],)
|
||||
|
||||
|
||||
def test_aggregate_stats_incremental_resume():
|
||||
"""Verify conservative bounds remain associative across incremental additions."""
|
||||
# Start with episode 1 stats (narrow distribution)
|
||||
ep1_stats = {
|
||||
"action": {
|
||||
"min": np.array([-10.0, -5.0]),
|
||||
"max": np.array([10.0, 5.0]),
|
||||
"mean": np.array([0.0, 0.0]),
|
||||
"std": np.array([3.0, 1.5]),
|
||||
"count": np.array([500]),
|
||||
"q01": np.array([-9.0, -4.5]),
|
||||
"q99": np.array([9.0, 4.5]),
|
||||
},
|
||||
}
|
||||
|
||||
# Episode 2: wider distribution on dim 0
|
||||
ep2_stats = {
|
||||
"action": {
|
||||
"min": np.array([-30.0, -5.0]),
|
||||
"max": np.array([40.0, 6.0]),
|
||||
"mean": np.array([5.0, 0.5]),
|
||||
"std": np.array([15.0, 2.0]),
|
||||
"count": np.array([100]),
|
||||
"q01": np.array([-25.0, -4.0]),
|
||||
"q99": np.array([35.0, 5.5]),
|
||||
},
|
||||
}
|
||||
|
||||
# First aggregation: ep1 + ep2 (simulates save_episode for ep2)
|
||||
cumulative = aggregate_stats([ep1_stats, ep2_stats])
|
||||
|
||||
# q01 should take min (conservative lower bound)
|
||||
np.testing.assert_allclose(cumulative["action"]["q01"], np.array([-25.0, -4.5]))
|
||||
# q99 should take max (conservative upper bound)
|
||||
np.testing.assert_allclose(cumulative["action"]["q99"], np.array([35.0, 5.5]))
|
||||
|
||||
# Episode 3: even wider on dim 1
|
||||
ep3_stats = {
|
||||
"action": {
|
||||
"min": np.array([-8.0, -20.0]),
|
||||
"max": np.array([8.0, 25.0]),
|
||||
"mean": np.array([0.0, 3.0]),
|
||||
"std": np.array([2.0, 8.0]),
|
||||
"count": np.array([50]),
|
||||
"q01": np.array([-7.0, -18.0]),
|
||||
"q99": np.array([7.0, 22.0]),
|
||||
},
|
||||
}
|
||||
|
||||
# Second aggregation: cumulative + ep3 (simulates save_episode for ep3)
|
||||
cumulative2 = aggregate_stats([cumulative, ep3_stats])
|
||||
|
||||
# Bounds should widen monotonically
|
||||
np.testing.assert_allclose(cumulative2["action"]["q01"], np.array([-25.0, -18.0]))
|
||||
np.testing.assert_allclose(cumulative2["action"]["q99"], np.array([35.0, 22.0]))
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user