mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-26 03:06:01 +00:00
Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Martino Russi
|
2cfd5971a8 | ||
|
Martino Russi
|
30cbea056e | ||
|
Martino Russi
|
b5201f6c15 |
@@ -55,7 +55,7 @@ jobs:
|
||||
github.repository == 'huggingface/lerobot'
|
||||
permissions:
|
||||
contents: read
|
||||
uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@e60a538eea9817ab312196d0d233604b01697265 # main
|
||||
uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391 # main
|
||||
with:
|
||||
commit_sha: ${{ github.sha }}
|
||||
package: lerobot
|
||||
@@ -78,7 +78,7 @@ jobs:
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@e60a538eea9817ab312196d0d233604b01697265 # main
|
||||
uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391 # main
|
||||
with:
|
||||
commit_sha: ${{ github.event.pull_request.head.sha }}
|
||||
pr_number: ${{ github.event.number }}
|
||||
|
||||
@@ -169,8 +169,6 @@
|
||||
- sections:
|
||||
- local: phone_teleop
|
||||
title: Phone
|
||||
- local: isaac_teleop
|
||||
title: Isaac Teleop
|
||||
title: "Teleoperators"
|
||||
- sections:
|
||||
- local: cameras
|
||||
|
||||
@@ -81,12 +81,6 @@ merged. Both prompts also carry a causal **event-boundary** definition (a
|
||||
new event starts when an object becomes held / is released / reaches a new
|
||||
location / a lid changes state / contents move) to sharpen where cuts land.
|
||||
|
||||
Optionally, a third **seeded-relabel** pass (`--plan.subtask_seeded_relabel`)
|
||||
revisits each span with its previous/current/next segment contact sheets and
|
||||
minimally corrects the label, using the first label as a prior — it keeps the
|
||||
boundaries fixed and only sharpens wording, at the cost of one extra call per
|
||||
subtask.
|
||||
|
||||
The resulting spans are then stitched into a gap-free, full-episode
|
||||
cover, so **every frame has exactly one active subtask**. See
|
||||
[`run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py)
|
||||
@@ -163,33 +157,30 @@ Every module is on by default and can be toggled independently (set to
|
||||
|
||||
### The VLM (`--vlm.*`)
|
||||
|
||||
| Flag | Default | What it does |
|
||||
| -------------------------- | ------------------ | ------------------------------------------------------------------------------------ |
|
||||
| `--vlm.model_id` | `Qwen/Qwen3.6-27B` | The model to serve and prompt. |
|
||||
| `--vlm.camera_key` | first `images.*` | Which camera every prompt is grounded on. |
|
||||
| `--vlm.serve_command` | auto | The exact `vllm serve …` command (set TP size, GPU memory, `--max-model-len` here). |
|
||||
| `--vlm.parallel_servers` | `1` | Independent servers for round-robin routing (one per GPU). |
|
||||
| `--vlm.num_gpus` | `0` | GPUs per server (`0` = one each). |
|
||||
| `--vlm.client_concurrency` | `16` | In-flight requests across all servers. |
|
||||
| `--vlm.max_new_tokens` | `512` | Generation cap per call. |
|
||||
| `--vlm.temperature` | `0.2` | Sampling temperature. |
|
||||
| `--vlm.reasoning_effort` | `null` | Thinking-budget hint (`low`/`medium`/`high`) forwarded to OpenAI-compatible servers. |
|
||||
| Flag | Default | What it does |
|
||||
| -------------------------- | ------------------ | ----------------------------------------------------------------------------------- |
|
||||
| `--vlm.model_id` | `Qwen/Qwen3.6-27B` | The model to serve and prompt. |
|
||||
| `--vlm.camera_key` | first `images.*` | Which camera every prompt is grounded on. |
|
||||
| `--vlm.serve_command` | auto | The exact `vllm serve …` command (set TP size, GPU memory, `--max-model-len` here). |
|
||||
| `--vlm.parallel_servers` | `1` | Independent servers for round-robin routing (one per GPU). |
|
||||
| `--vlm.num_gpus` | `0` | GPUs per server (`0` = one each). |
|
||||
| `--vlm.client_concurrency` | `16` | In-flight requests across all servers. |
|
||||
| `--vlm.max_new_tokens` | `512` | Generation cap per call. |
|
||||
| `--vlm.temperature` | `0.2` | Sampling temperature. |
|
||||
|
||||
### Subtasks / plan / memory (`--plan.*`)
|
||||
|
||||
| Flag | Default | What it does |
|
||||
| ------------------------------- | ---------- | ---------------------------------------------------------------------------------------------------------------------------- |
|
||||
| `--plan.frames_per_second` | `2.0` | Frame sampling rate for the contact sheets (`2.0` = one frame every 0.5s). |
|
||||
| `--plan.max_frames_per_prompt` | `60` | Frame budget per VLM call. Episodes whose sampling exceeds this are auto-windowed at the same density, then stitched. |
|
||||
| `--plan.contact_sheet_columns` | `5` | Columns per contact-sheet grid (`contact_sheet_frames_per_sheet` tiles, time row-major). |
|
||||
| `--plan.plan_max_steps` | `8` | Upper bound on subtasks per episode. |
|
||||
| `--plan.subtask_describe_first` | `true` | Run the describe→segment grounding pass (best subtask quality; +1 call/episode). |
|
||||
| `--plan.subtask_seeded_relabel` | `false` | Second pass: re-label each subtask from its prev/current/next contact sheets, seeded with the first label (+1 call/subtask). |
|
||||
| `--plan.subtask_relabel_frames` | `5` | Frames sampled uniformly per segment sheet in the relabel pass (only used when `subtask_seeded_relabel=true`). |
|
||||
| `--plan.emit_plan` | `true` | Emit the numbered `plan` rows (`false` = subtasks + memory only). |
|
||||
| `--plan.emit_memory` | `true` | Emit the `memory` rows (`false` = subtasks + plan only); symmetric to `emit_plan`. |
|
||||
| `--plan.n_task_rephrasings` | `10` | How many `task_aug` rephrasings to emit (`0` disables). |
|
||||
| `--plan.derive_task_from_video` | `if_short` | Use the dataset task as-is (`off`), only when it's missing/short (`if_short`), or always re-derive from video (`always`). |
|
||||
| Flag | Default | What it does |
|
||||
| ------------------------------- | ---------- | ------------------------------------------------------------------------------------------------------------------------- |
|
||||
| `--plan.frames_per_second` | `2.0` | Frame sampling rate for the contact sheets (`2.0` = one frame every 0.5s). |
|
||||
| `--plan.max_frames_per_prompt` | `60` | Frame budget per VLM call. Episodes whose sampling exceeds this are auto-windowed at the same density, then stitched. |
|
||||
| `--plan.contact_sheet_columns` | `5` | Columns per contact-sheet grid (`contact_sheet_frames_per_sheet` tiles, time row-major). |
|
||||
| `--plan.plan_max_steps` | `8` | Upper bound on subtasks per episode. |
|
||||
| `--plan.subtask_describe_first` | `true` | Run the describe→segment grounding pass (best subtask quality; +1 call/episode). |
|
||||
| `--plan.emit_plan` | `true` | Emit the numbered `plan` rows (`false` = subtasks + memory only). |
|
||||
| `--plan.emit_memory` | `true` | Emit the `memory` rows (`false` = subtasks + plan only); symmetric to `emit_plan`. |
|
||||
| `--plan.n_task_rephrasings` | `10` | How many `task_aug` rephrasings to emit (`0` disables). |
|
||||
| `--plan.derive_task_from_video` | `if_short` | Use the dataset task as-is (`off`), only when it's missing/short (`if_short`), or always re-derive from video (`always`). |
|
||||
|
||||
### Interjections + VQA
|
||||
|
||||
|
||||
@@ -150,14 +150,14 @@ class MyPolicy(PreTrainedPolicy):
|
||||
|
||||
The methods called by the train/eval loops:
|
||||
|
||||
| Method | Used by | What it does |
|
||||
| ----------------------------------------------------------------- | ----------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| `reset() -> None` | `lerobot-eval` | Clear per-episode state at the start of each episode. |
|
||||
| `select_action(batch, **kwargs) -> Tensor` | `lerobot-eval` | Return the next action `(B, action_dim)`. Called every step. |
|
||||
| `predict_action_chunk(batch, **kwargs) -> Tensor` | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk. |
|
||||
| `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train` | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting. |
|
||||
| `get_optim_params() -> dict` | the optimizer | Return `self.parameters()` for simple policies; return a named parameter dict for multi-optimizer policies (see `get_optim_params` in [`modeling_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/modeling_act.py) for a per-group learning-rate example). |
|
||||
| `update() -> None` _(optional)_ | `lerobot-train` | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this). |
|
||||
| Method | Used by | What it does |
|
||||
| ----------------------------------------------------------------- | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| `reset() -> None` | `lerobot-eval` | Clear per-episode state at the start of each episode. |
|
||||
| `select_action(batch, **kwargs) -> Tensor` | `lerobot-eval` | Return the next action `(B, action_dim)`. Called every step. |
|
||||
| `predict_action_chunk(batch, **kwargs) -> Tensor` | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk. |
|
||||
| `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train` | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting. |
|
||||
| `get_optim_params() -> dict` | the optimizer | Return `self.parameters()` for simple policies; return a named parameter dict for [multi-optimizer policies](https://github.com/huggingface/lerobot/blob/ecd38c50d7d15b4184cf42649ff1185ee2e11eeb/src/lerobot/policies/sac/modeling_sac.py#L61-L73). |
|
||||
| `update() -> None` _(optional)_ | `lerobot-train` | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this). |
|
||||
|
||||
Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constants`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/utils/constants.py): `OBS_STATE` (`observation.state.<motor>`), `OBS_IMAGES` (`observation.images.<camera>`), `OBS_LANGUAGE`, `ACTION`, etc. Reuse the constants — don't invent new prefixes.
|
||||
|
||||
@@ -295,10 +295,12 @@ The file names are load-bearing: the factory does lazy imports by name, and the
|
||||
|
||||
### Wiring
|
||||
|
||||
Two places need to know about your policy. All by name.
|
||||
Four places need to know about your policy. All by name.
|
||||
|
||||
1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. This import is what registers your policy: `@PreTrainedConfig.register_subclass("my_policy")` runs, and from then on the factory resolves everything by convention. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast).
|
||||
2. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page.
|
||||
1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast).
|
||||
2. **`factory.py:get_policy_class`** — add a branch returning `MyPolicy` from a lazy import.
|
||||
3. **`factory.py:make_policy_config`** and **`factory.py:make_pre_post_processors`** — same idea, two more branches.
|
||||
4. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page.
|
||||
|
||||
Mirror an existing policy that's structurally similar to yours; the diff is small.
|
||||
|
||||
@@ -330,10 +332,6 @@ This way:
|
||||
|
||||
Add a matching extra to [`pyproject.toml`](https://github.com/huggingface/lerobot/blob/main/pyproject.toml) `[project.optional-dependencies]` and include it in the `all` extra so `pip install 'lerobot[all]'` keeps installing everything.
|
||||
|
||||
### Avoid copying a modeling file — subclass it
|
||||
|
||||
If your policy needs to modify a backbone that already exists in `transformers` (custom conditioning, extra inputs, a swapped sub-module), **do not vendor a copy of its `modeling_*.py`**. Instead, subclass the smallest upstream unit and override only what changes. [`pi_gemma.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi_gemma.py) is the canonical reference: it injects AdaRMS conditioning into PaliGemma/Gemma in ~370 lines by subclassing `GemmaModel`/`PaliGemmaModel` and overriding the decoder-layer forward, instead of forking the ~2,000-line modeling file. Model surgery on a _loaded_ native model is also fine (layer truncation, tokenizer expansion, hidden-state capture — see `evo1/internvl3_embedder.py`, `eo1/modeling_eo1.py`, `groot/groot_n1_7.py` for working examples). Reviewers will ask for this pattern when a PR arrives with a copied modeling file; the only accepted exception is a model that does not exist in `transformers` at all.
|
||||
|
||||
### Benchmarks and a published checkpoint
|
||||
|
||||
A new policy is much easier to review — and far more useful — when it ships with a working checkpoint and at least one number you can reproduce.
|
||||
@@ -369,7 +367,7 @@ If your policy is real-robot-only and no sim benchmark applies, swap the sim eva
|
||||
The general expectations are in [`CONTRIBUTING.md`](https://github.com/huggingface/lerobot/blob/main/CONTRIBUTING.md) and the [PR template](https://github.com/huggingface/lerobot/blob/main/.github/PULL_REQUEST_TEMPLATE.md). On top of those, reviewers will look for:
|
||||
|
||||
- [ ] `MyPolicy` and `MyPolicyConfig` cover the surface above; `__init_subclass__` accepts the class.
|
||||
- [ ] `policies/__init__.py` re-exports the config (this registers the policy; the factory resolves modeling/processor by naming convention).
|
||||
- [ ] `factory.py` and `policies/__init__.py` are wired (lazy imports for modeling).
|
||||
- [ ] `make_my_policy_pre_post_processors` follows the naming convention.
|
||||
- [ ] Optional deps live behind a `[project.optional-dependencies]` extra and the `TYPE_CHECKING + require_package` guard.
|
||||
- [ ] `tests/policies/` updated; backward-compat artifact committed & policy-specific tests.
|
||||
|
||||
@@ -162,11 +162,11 @@ Preliminary LeRobot integration results (GR00T-LeRobot, `eval.n_episodes >= 50`
|
||||
|
||||
| Suite | Success rate | Checkpoint |
|
||||
| ---------------- | -----------: | ------------------------------------------------------------------------------------------------------------- |
|
||||
| LIBERO Spatial | 95% | [nvidia/gr00t17-lerobot-libero_spatial-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_spatial-640) |
|
||||
| LIBERO Object | 100% | [nvidia/gr00t17-lerobot-libero_object-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_object-640) |
|
||||
| LIBERO Goal | 98% | [nvidia/gr00t17-lerobot-libero_goal-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_goal-640) |
|
||||
| LIBERO 10 (Long) | 93% | [nvidia/gr00t17-lerobot-libero_10-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_10-640) |
|
||||
| **Average** | **96.5%** | |
|
||||
| LIBERO Spatial | 91% | [nvidia/gr00t17-lerobot-libero_spatial-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_spatial-640) |
|
||||
| LIBERO Object | 81% | [nvidia/gr00t17-lerobot-libero_object-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_object-640) |
|
||||
| LIBERO Goal | 97% | [nvidia/gr00t17-lerobot-libero_goal-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_goal-640) |
|
||||
| LIBERO 10 (Long) | 84% | [nvidia/gr00t17-lerobot-libero_10-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_10-640) |
|
||||
| **Average** | **88.25%** | |
|
||||
|
||||
```bash
|
||||
export MODEL_ID=your_trained_model_on_huggingface
|
||||
|
||||
@@ -1,397 +0,0 @@
|
||||
# Isaac Teleop
|
||||
|
||||
Control your robot with NVIDIA [Isaac Teleop](https://github.com/NVIDIA/IsaacTeleop), a
|
||||
multi-modal teleoperation framework. Isaac Teleop drives a single `TeleopSession` from a range
|
||||
of input devices — XR (VR) controllers, hand tracking, full-body tracking, Manus gloves, foot
|
||||
pedals, and more.
|
||||
|
||||
In LeRobot, Isaac Teleop ships as a self-contained example under
|
||||
[`examples/isaac_teleop_to_so101/`](https://github.com/huggingface/lerobot/tree/main/examples/isaac_teleop_to_so101).
|
||||
Each Isaac Teleop input device is its own `Teleoperator` subclass in the example's
|
||||
`isaac_teleop` package, sharing one session lifecycle (see `IsaacTeleopTeleoperator`). The
|
||||
devices available today are the **XR controller** (`XRController`) and a back-drivable
|
||||
**SO-101 leader arm** (`SO101LeaderArm`); Manus gloves and hand/full-body tracking are the
|
||||
natural next devices. This guide focuses on the XR controller; the SO-101 leader is summarized
|
||||
under [Run the example](#step-3-run-the-example).
|
||||
|
||||
**In this guide you'll learn:**
|
||||
|
||||
- How an Isaac Teleop device drives a robot end‑effector (EE) target
|
||||
- How the _clutch_ (squeeze/grip on the XR controller) engages teleoperation without jerking the arm
|
||||
- How to run the SO‑101 teleoperation example and tune motion / gripper / IK
|
||||
|
||||
## Installation
|
||||
|
||||
The example lives in the LeRobot repository (it is not part of the `lerobot` pip package), so
|
||||
clone the repo and install from source. The canonical, always-up-to-date install and usage
|
||||
reference is the example's
|
||||
[`README.md`](https://github.com/huggingface/lerobot/tree/main/examples/isaac_teleop_to_so101/README.md);
|
||||
in short:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/huggingface/lerobot.git
|
||||
cd lerobot
|
||||
uv pip install -e ".[feetech,kinematics,dataset]" "huggingface_hub>=1.5"
|
||||
uv pip install "isaacteleop[cloudxr,retargeters-lite]~=1.3.131" "scipy>=1.14"
|
||||
```
|
||||
|
||||
`isaacteleop` is published on public PyPI (Linux only). The `cloudxr` extra brings the CloudXR
|
||||
runtime bindings; `retargeters-lite` is the scipy-based retargeter path that resolves on both
|
||||
x86_64 and ARM (on aarch64 — e.g. a DGX Spark — the full `retargeters` extra does not resolve
|
||||
because of its `dex-retargeting`/`nlopt` pins, which is why it is not the default here). On
|
||||
x86_64 you can additionally install the full retargeter stack:
|
||||
|
||||
```bash
|
||||
uv pip install "isaacteleop[retargeters]~=1.3.131"
|
||||
```
|
||||
|
||||
### Set up CloudXR and connect a headset
|
||||
|
||||
Isaac Teleop streams the headset to your machine over **NVIDIA CloudXR**, which provides the
|
||||
OpenXR runtime the session connects to. By default LeTeleop **auto-launches the CloudXR runtime
|
||||
for you** when you call `teleop_device.connect()` — you no longer have to run `python -m
|
||||
isaacteleop.cloudxr` and `source cloudxr.env` in a separate shell. All you need is a supported
|
||||
headset connected and the CloudXR firewall ports open. Follow the Isaac Teleop
|
||||
[Quick Start](https://nvidia.github.io/IsaacTeleop/main/getting_started/quick_start.html) for the
|
||||
headset-pairing and firewall details.
|
||||
|
||||
**First run (EULA).** The very first launch must accept the NVIDIA CloudXR EULA. The auto-launch
|
||||
prompts for it **on stdin**, so on a headless machine it will hang waiting for input. Bootstrap
|
||||
the EULA once, interactively, with:
|
||||
|
||||
```bash
|
||||
python -m isaacteleop.cloudxr --accept-eula # one-time: accept the CloudXR EULA
|
||||
```
|
||||
|
||||
After that, `connect()` launches the runtime non-interactively. The launch **blocks for ~30s**
|
||||
while the runtime comes up.
|
||||
|
||||
**Configuration.** Two fields on `IsaacTeleopConfig` (shared by every device) control this:
|
||||
|
||||
- `auto_launch_cloudxr` (default `True`) — whether `connect()` starts the runtime. Set `False`
|
||||
when CloudXR is already running externally.
|
||||
- `cloudxr_env_file` (default `None`) — an optional CloudXR device-profile `.env` selecting the
|
||||
headset transport (e.g. an Apple Vision Pro profile). This is launcher **input**; it is not the
|
||||
`~/.cloudxr/run/cloudxr.env` **output** file the old manual flow told you to `source`. `None`
|
||||
keeps the default auto-WebRTC profile — though the SO-101 example overrides it to the
|
||||
`default.env` shipped next to `teleoperate.py` unless you pass `--teleop.cloudxr_env_file`.
|
||||
|
||||
**Opting out.** To skip the auto-launch (CloudXR already running), either set
|
||||
`auto_launch_cloudxr=False` or export:
|
||||
|
||||
```bash
|
||||
export LEROBOT_CLOUDXR_SKIP_AUTOLAUNCH=1
|
||||
```
|
||||
|
||||
The **env var takes precedence over the config field**: if `LEROBOT_CLOUDXR_SKIP_AUTOLAUNCH=1` is
|
||||
set, the auto-launch is skipped even when `auto_launch_cloudxr=True`. This variable is
|
||||
**independent** of Isaac Lab's `ISAACLAB_CXR_SKIP_AUTOLAUNCH` — setting one does not affect the
|
||||
other.
|
||||
|
||||
**One teleoperator per process.** The CloudXR runtime configures the environment process-wide (a
|
||||
singleton), so run a single Isaac Teleop teleoperator per process.
|
||||
|
||||
**Shutting down.** Always call `teleop_device.disconnect()` on exit — including on Ctrl-C. Wrap
|
||||
your teleoperation loop in `try/finally` and call `disconnect()` in the `finally`. This tears down
|
||||
the OpenXR session **before** the CloudXR runtime, which is the required order; the launcher's
|
||||
`atexit` hook only reaps the runtime and does not run the session's `__exit__`, so without an
|
||||
explicit `disconnect()` an interrupted run shuts down in the wrong order.
|
||||
|
||||
```python
|
||||
teleop_device.connect()
|
||||
try:
|
||||
while True:
|
||||
action = teleop_device.get_action()
|
||||
# ... drive the robot ...
|
||||
finally:
|
||||
teleop_device.disconnect()
|
||||
```
|
||||
|
||||
See [System Requirements](https://nvidia.github.io/IsaacTeleop/main/references/requirements.html)
|
||||
for supported OS / GPU / CloudXR versions and headsets.
|
||||
|
||||
## How it works
|
||||
|
||||
The XR controller is one Isaac Teleop **input** device. `XRController` is a deliberately thin
|
||||
reader: it exposes the **raw** controller grip pose — already statically rebased into the robot
|
||||
base frame — plus the squeeze and trigger analog values. It has **no** retargeters and **no**
|
||||
clutch logic of its own. The clutch (engage latch + delta rebasing onto the EE) and the gripper
|
||||
mapping live downstream in the example loop, which then feeds LeRobot's existing closed‑loop
|
||||
Cartesian IK pipeline — the same one the phone teleoperator uses. The device‑specific pieces are
|
||||
`XRController`, the loop's `Clutch`, and `MapXRControllerActionToRobotAction`; everything downstream
|
||||
(`EEBoundsAndSafety`, `InverseKinematicsEEToJoints`) is shared, and a future device (e.g. Manus
|
||||
gloves) would swap in its own `teleop_<device>.py` + processor while reusing the rest.
|
||||
|
||||
`XRController._build_pipeline` wires Isaac Teleop's `ControllersSource` — statically rebased into
|
||||
the robot base frame by the native `ControllerTransform` (`base_T_anchor`) — and exposes the
|
||||
transformed controller stream verbatim. `get_action()` reads the grip pose, squeeze, and trigger
|
||||
straight off it; the session is always stepped `RUNNING` (there is no clutch retargeter to gate).
|
||||
|
||||
The `Clutch` class (in `examples/isaac_teleop_to_so101/isaac_teleop/clutch.py`, driven by the
|
||||
loop in `common.py`) mirrors Isaac Teleop's `SO101ClutchRetargeter`, but lives in-loop so the
|
||||
device can stay a thin reader:
|
||||
|
||||
- It latches its engage origin on the squeeze **engage edge** (the frame the squeeze first crosses
|
||||
`clutch_threshold`) and rebases both position and orientation around it, so engaging does not
|
||||
teleport the arm. `Clutch.rebase` returns the absolute base-frame target as a `(pos, quat)`
|
||||
pair, which the loop concatenates into the 7D `ee_pose` fed to the processor.
|
||||
- The analog trigger becomes a gripper `closedness` in `[0, 1]` (0 = open, 1 = closed),
|
||||
proportional to the trigger pull, which `MapXRControllerActionToRobotAction` maps to a jaw target.
|
||||
|
||||
See the Isaac Teleop
|
||||
[Retargeting interface](https://nvidia.github.io/IsaacTeleop/main/references/retargeting/index.html)
|
||||
and [architecture overview](https://nvidia.github.io/IsaacTeleop/main/overview/architecture.html)
|
||||
for how source nodes and retargeters compose.
|
||||
|
||||
```text
|
||||
VR controller (OpenXR)
|
||||
│
|
||||
▼
|
||||
XRController.get_action() ── raw base-frame grip_pos / grip_quat + squeeze + trigger
|
||||
│ (TeleopSession always stepped RUNNING; clutch lives downstream)
|
||||
▼
|
||||
Clutch.rebase(grip_pos, grip_quat) ── engage-relative delta applied to the EE home (pos + orient)
|
||||
│ ee_pose (7) / closedness → absolute ee_pose; closedness = trigger
|
||||
▼
|
||||
MapXRControllerActionToRobotAction ── absolute ee.x/y/z; ee.w* = orientation rotvec target;
|
||||
│ ee.x/y/z / ee.w* / ee.gripper_pos ee.gripper_pos = (1 - closedness) * 100
|
||||
▼
|
||||
EEBoundsAndSafety ── workspace clip + per-frame step clamp (clamp+warn)
|
||||
│
|
||||
▼
|
||||
InverseKinematicsEEToJoints ── closed-loop Placo IK; position + soft-orientation
|
||||
│ (orientation_weight=0.01) (passes ee.gripper_pos → gripper.pos)
|
||||
▼
|
||||
SO-101 follower joint targets
|
||||
```
|
||||
|
||||
### The clutch: owned by the example loop
|
||||
|
||||
Unlike the phone pipeline (which splits the clutch across `MapPhoneActionToRobotAction` and
|
||||
`EEReferenceAndDelta`), the XR clutch lives entirely in the example loop's `Clutch` class. It emits
|
||||
an **absolute** EE pose, so there is no `EEReferenceAndDelta` stage and no delta accumulation in the
|
||||
processor — `MapXRControllerActionToRobotAction` is a pure, stateless per‑frame mapping.
|
||||
|
||||
The clutch latches its engage origin on the squeeze **engage edge** (the moment the squeeze crosses
|
||||
`clutch_threshold`) and drives the EE from the motion _relative_ to that origin, so the arm does not
|
||||
teleport on engage. On **every** engage — startup and mid‑task re‑clutch alike — the home
|
||||
_position_ is latched from forward kinematics on the arm's **measured joints**, so the home equals
|
||||
where the arm physically is even if it moved while disengaged, and the engage is jump‑free. The
|
||||
home _orientation_ keeps the last commanded rotation: the 5‑DOF arm tracks orientation only
|
||||
softly, so latching the measured wrist orientation would inject its tracking offset into the
|
||||
command on every re‑clutch.
|
||||
|
||||
## Controls
|
||||
|
||||
- **Squeeze / grip** — the **clutch** (deadman). Hold it past `clutch_threshold` to engage
|
||||
teleoperation; release to pause. Each engage re‑captures the origin, so you can reposition
|
||||
your hand while paused and re‑engage without the arm jumping (index/clutch style).
|
||||
- **Trigger** — the **gripper**, controlled **analog**. The jaw tracks the trigger
|
||||
proportionally — a half‑pressed trigger leaves the jaw half‑closed — via a closedness in
|
||||
`[0, 1]` (0 = open, 1 = closed) that maps to an absolute gripper joint target.
|
||||
- **Controller orientation** — the **wrist**. The clutch rebases the controller orientation
|
||||
(engage‑relative, base‑frame) into a soft IK orientation target the wrist tracks alongside
|
||||
position. On the 5‑DOF SO‑101 the wrist follows the hand only partially by design — see
|
||||
`orientation_weight` below.
|
||||
|
||||
## Get started
|
||||
|
||||
### Step 1: Create the teleoperator
|
||||
|
||||
```python
|
||||
# Run from the repo root so the `examples` package is importable.
|
||||
from examples.isaac_teleop_to_so101.isaac_teleop import XRController, XRControllerConfig
|
||||
|
||||
teleop_config = XRControllerConfig(
|
||||
hand_side="right", # "left" or "right" controller
|
||||
clutch_threshold=0.5, # squeeze value above which the clutch engages
|
||||
)
|
||||
teleop_device = XRController(teleop_config)
|
||||
```
|
||||
|
||||
`XRController.get_action()` returns the **raw** base‑frame controller pose, not a clutch‑rebased
|
||||
target: `grip_pos` (3,) `[x, y, z]` [m] and `grip_quat` (4,) `[qx, qy, qz, qw]` in the robot base
|
||||
frame, plus scalar `squeeze` and `trigger` analog values in `[0, 1]`. The example loop's `Clutch`
|
||||
turns these into the absolute `ee_pose`, and the squeeze is thresholded by the loop against
|
||||
`clutch_threshold` to engage.
|
||||
|
||||
### Step 2: Connect
|
||||
|
||||
Calling `teleop_device.connect()` first auto-launches the CloudXR runtime (unless you opted out —
|
||||
see [Set up CloudXR and connect a headset](#set-up-cloudxr-and-connect-a-headset); this blocks for
|
||||
~30s and on the first run prompts for the EULA on stdin), then starts the Isaac Teleop
|
||||
[`TeleopSession`](https://nvidia.github.io/IsaacTeleop/main/getting_started/teleop_session.html)
|
||||
(opens the OpenXR session and discovers the controllers). XR controllers are self‑calibrating, so
|
||||
there is no manual calibration step — the clutch handles re‑centering each time you engage. Pair
|
||||
`connect()` with a `try/finally` that calls `disconnect()` so the session tears down before the
|
||||
runtime on exit/Ctrl-C.
|
||||
|
||||
### Step 3: Run the example
|
||||
|
||||
The example assumes you configured your robot (SO‑101 follower) and set the correct serial port.
|
||||
|
||||
The **robot URDF and its meshes are fetched automatically** on first run: the XR device downloads
|
||||
the SO-101 URDF from the
|
||||
[`lerobot/robot-urdfs` Hugging Face bucket](https://huggingface.co/buckets/lerobot/robot-urdfs/tree/so101)
|
||||
into the LeRobot cache (`HF_LEROBOT_HOME/robot-urdfs/so101/`) and reuses it after, so there is no
|
||||
separate download step :
|
||||
|
||||
```bash
|
||||
python -m examples.isaac_teleop_to_so101.teleoperate --robot.type=so101_follower --robot.port=/dev/ttyACM0 \
|
||||
--robot.id=so101_follower_arm --teleop.type=xr_controller
|
||||
```
|
||||
|
||||
The CLI is `lerobot-teleoperate`-style (draccus): `--robot.*` configures the SO-101 follower and
|
||||
`--teleop.type` selects the Isaac input device (`xr_controller` | `so101_leader`), with
|
||||
`--teleop.*` its device knobs. `--teleop.type=xr_controller` runs the XR-controller path described
|
||||
above. The startup safety contract: by default it slews all joints to a default reset pose over
|
||||
`--reset_duration` seconds (`--reset_to_origin=false` keeps the arm where it is), then seeds the
|
||||
clutch home from the arm's measured pose so the first engage is jump-free; the follower is
|
||||
commanded only while the clutch is engaged.
|
||||
|
||||
**Customizing the reset pose.** The reset pose ships as a built-in default (a comfortable mid-range
|
||||
pose) and works out of the box — you do **not** need to record anything. To tailor it to your setup,
|
||||
back-drive the arm to the pose you want and run
|
||||
`python -m examples.isaac_teleop_to_so101.override_reset_pose --id <robot.id>`; it writes the
|
||||
current joints to a per-arm file in the LeRobot cache
|
||||
(`HF_LEROBOT_HOME/reset_poses/<robot.name>/<robot.id>.json`, keyed like calibration), which then takes
|
||||
priority over the built-in default on the next run. Because it lives in the user-local cache (not
|
||||
the repo), your override stays on your machine, and both `teleoperate` and `record` honor it
|
||||
when launched with the same `--robot.id`.
|
||||
|
||||
The other device, `--teleop.type=so101_leader`, mirrors the follower 1:1 from a back-drivable
|
||||
SO-101 _leader arm_ whose joints are streamed by Isaac Teleop's native `so101_leader` plugin (no
|
||||
clutch, no IK — the leader and follower share the SO-101 kinematics).
|
||||
|
||||
The `so101_leader_plugin` binary is a C++ plugin that is **not** part of the `isaacteleop` pip
|
||||
package — you build it from the Isaac Teleop source tree. Follow
|
||||
[Build Isaac Teleop from source](https://nvidia.github.io/IsaacTeleop/main/getting_started/build_from_source/index.html)
|
||||
(in short, from your Isaac Teleop checkout: `cmake -B build && cmake --build build --parallel &&
|
||||
cmake --install build`); the build installs the plugins under `<IsaacTeleop>/install/plugins/`, so
|
||||
the binary lands at `install/plugins/so101_leader/so101_leader_plugin` — the `--launch_plugin` path
|
||||
below. See the plugin's own `README.md` (next to the binary) for its serial/calibration details.
|
||||
|
||||
Point `--teleop.port` at the physical leader's serial port and `--launch_plugin` at that plugin
|
||||
binary to have the script spawn it after CloudXR is up:
|
||||
|
||||
```bash
|
||||
python -m examples.isaac_teleop_to_so101.teleoperate --robot.type=so101_follower --robot.port=/dev/ttyACM0 \
|
||||
--robot.id=so101_follower_arm --teleop.type=so101_leader \
|
||||
--teleop.port=/dev/ttyACM1 --teleop.id=so101_leader_arm \
|
||||
--launch_plugin=/code/Teleop/install/plugins/so101_leader/so101_leader_plugin
|
||||
```
|
||||
|
||||
(Note `so101_leader` here is the _Isaac_ leader, resolved against the Isaac Teleop device
|
||||
registry, distinct from `lerobot-teleoperate`'s serial `so101_leader`.) When a `--teleop.port` is
|
||||
set, the plugin's tick→radian calibration is inferred from `--teleop.id` and passed to the plugin
|
||||
as its third positional arg — the LeRobot-format JSON at
|
||||
`HF_LEROBOT_CALIBRATION/teleoperators/so_leader/<id>.json`, the same file the serial SO-101 leader
|
||||
uses (`lerobot-calibrate --teleop.type=so101_leader --teleop.id=<id>`). If it is missing the script
|
||||
warns and the plugin uses built-in defaults. Run `python -m examples.isaac_teleop_to_so101.teleoperate --help` for all flags. Its
|
||||
startup safety contract: by default the follower is
|
||||
slewed to the leader's first reading over `--align_duration` seconds (`--align=false` to skip) so
|
||||
the arm does not snap when the mirror begins, and while the leader stream is stale the follower is
|
||||
held at its measured pose.
|
||||
|
||||
The URDF fetch uses `huggingface_hub` (already a LeRobot dependency) against the public
|
||||
`lerobot/robot-urdfs` bucket, so it needs no login. It is cached under
|
||||
`HF_LEROBOT_HOME/robot-urdfs/so101/`; delete that folder to force a re‑download.
|
||||
|
||||
Then, in your headset: squeeze and hold the grip to engage, move the controller to drive the
|
||||
arm, twist/tilt it to orient the wrist, and press the trigger to close the gripper
|
||||
(proportionally — release to open).
|
||||
|
||||
To record a dataset (not just teleoperate), use `record.py` in the same folder. It dispatches on
|
||||
`--teleop.type` (`xr_controller` | `so101_leader`) exactly like `teleoperate.py`, so either device
|
||||
can drive the follower, and it saves the commanded joints to a LeRobot dataset (`lerobot-record`-style
|
||||
`--dataset.*` flags). See its module docstring for the full CLI and the keyboard recording shortcuts.
|
||||
|
||||
## Important pipeline steps and options
|
||||
|
||||
The clutch already produces an absolute base‑frame pose, so the processor side is a thin
|
||||
**absolute‑pose** path — there is no frame remap, no delta accumulation, and no
|
||||
`EEReferenceAndDelta` stage.
|
||||
|
||||
- `MapXRControllerActionToRobotAction` is a stateless per‑frame mapping from the device output to
|
||||
the IK input contract. It writes the absolute base‑frame position, encodes the absolute
|
||||
orientation as a rotvec target, and inverts the closedness into a motor gripper target:
|
||||
|
||||
```python
|
||||
action["ee.x"], action["ee.y"], action["ee.z"] = ee_pose[:3] # absolute, base frame [m]
|
||||
action["ee.wx"], action["ee.wy"], action["ee.wz"] = orient_rotvec # orientation target (rotvec)
|
||||
action["ee.gripper_pos"] = (1 - closedness) * 100 # motor units; SO-101 calibrates 100 = open
|
||||
```
|
||||
|
||||
The gripper polarity (`100 = open, 0 = closed`) is a hardware‑calibration convention in the source — flip it there if the jaw opens when it should close.
|
||||
|
||||
- `EEBoundsAndSafety` clamps the EE to a workspace and rate‑limits per‑frame jumps. The clutch's
|
||||
no‑teleport keeps frames small, so `max_ee_step_m` mostly catches transient controller tracking
|
||||
glitches. The z floor is `0.0` (the table plane) so a stray target cannot drive the EE below the
|
||||
table; x/y stay at the loose `[-1, 1]` m box. Set `raise_on_jump=False` so an over‑limit frame is
|
||||
**clamped and warned** instead of raising — a crash mid‑loop would leave the arm uncontrolled:
|
||||
|
||||
```python
|
||||
EEBoundsAndSafety(
|
||||
end_effector_bounds={"min": [-1.0, -1.0, 0.0], "max": [1.0, 1.0, 1.0]},
|
||||
max_ee_step_m=0.10,
|
||||
raise_on_jump=False,
|
||||
)
|
||||
```
|
||||
|
||||
- `InverseKinematicsEEToJoints(initial_guess_current_joints=False, orientation_weight=0.01)` solves
|
||||
closed‑loop Placo IK. SO‑101 is a 5‑DOF arm, so the IK is position‑dominant; the small
|
||||
`orientation_weight` lets it softly track the orientation target carried in `ee.w*` so the wrist
|
||||
follows the hand, while the under‑determined roll stays partial by design. There is **no**
|
||||
`GripperVelocityToJoint`: the absolute `ee.gripper_pos` is passed straight to `gripper.pos`.
|
||||
`initial_guess_current_joints=False` warm‑starts each solve from the **previous IK solution**
|
||||
rather than re‑seeding from the measured joints, so the joint trajectory stays continuous
|
||||
frame‑to‑frame. Tune `orientation_weight` on hardware — too high fights position tracking, too
|
||||
low ignores the orientation command.
|
||||
|
||||
The example also gates safety at the loop level: after the startup reset slew (on by default —
|
||||
pass `--reset_to_origin=false` to keep the arm where it is), it commands the robot **only while
|
||||
the clutch is engaged**, and re‑sends the measured joints while disengaged, so releasing the
|
||||
clutch freezes the arm in place.
|
||||
|
||||
See the [Processors for Robots and Teleoperators](./processors_robots_teleop) guide for more on
|
||||
adapting the pipeline to other robots.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- **`ModuleNotFoundError: isaacteleop`** — the `isaacteleop` package is not installed in the
|
||||
active environment. Re-run the install command at the top of this guide:
|
||||
`uv pip install "isaacteleop[cloudxr,retargeters-lite]~=1.3.131"`.
|
||||
- **No controllers found** — make sure the CloudXR runtime is running, the firewall ports are
|
||||
whitelisted, and the headset is connected (see
|
||||
[Set up CloudXR and connect a headset](#set-up-cloudxr-and-connect-a-headset) and the Isaac
|
||||
Teleop [Quick Start](https://nvidia.github.io/IsaacTeleop/main/getting_started/quick_start.html)).
|
||||
- **CloudXR auto-launch failed** — `connect()` raises a `RuntimeError` if the runtime does not
|
||||
come up within its startup timeout. Check the launcher logs under `~/.cloudxr/logs`. Common
|
||||
causes: the EULA was never accepted (run `python -m isaacteleop.cloudxr --accept-eula` once,
|
||||
interactively — the auto-launch prompts on stdin and hangs headless), or the runtime is already
|
||||
running externally (set `LEROBOT_CLOUDXR_SKIP_AUTOLAUNCH=1` or `auto_launch_cloudxr=False` to
|
||||
skip the auto-launch).
|
||||
- **Arm does not move** — the clutch is a deadman: you must hold the squeeze/grip past
|
||||
`clutch_threshold`. Lower the threshold if your controller's squeeze is reported softly.
|
||||
- **Motion feels misaligned** — confirm the headset/play space orientation. The controller stream
|
||||
is rebased into the robot base frame by the `base_T_anchor` transform on `XRControllerConfig`
|
||||
(default: standard OpenXR → robot axis convention); adjust it if your anchor frame differs.
|
||||
|
||||
## Learn more
|
||||
|
||||
NVIDIA Isaac Teleop documentation ([docs home](https://nvidia.github.io/IsaacTeleop/),
|
||||
[GitHub](https://github.com/NVIDIA/IsaacTeleop)):
|
||||
|
||||
- [Quick Start](https://nvidia.github.io/IsaacTeleop/main/getting_started/quick_start.html) —
|
||||
install, run the CloudXR server, connect a headset, run a teleop example.
|
||||
- [TeleopSession](https://nvidia.github.io/IsaacTeleop/main/getting_started/teleop_session.html) —
|
||||
the session API `XRController` wraps.
|
||||
- [Retargeting interface](https://nvidia.github.io/IsaacTeleop/main/references/retargeting/index.html)
|
||||
and [architecture overview](https://nvidia.github.io/IsaacTeleop/main/overview/architecture.html) —
|
||||
how source nodes and retargeters compose into a pipeline.
|
||||
- [Build from source](https://nvidia.github.io/IsaacTeleop/main/getting_started/build_from_source/index.html) —
|
||||
build `isaacteleop` (and its C++ plugins, including the `so101_leader` plugin used above) from a
|
||||
local checkout.
|
||||
- [System Requirements](https://nvidia.github.io/IsaacTeleop/main/references/requirements.html) and
|
||||
the [CloudXR SDK docs](https://docs.nvidia.com/cloudxr-sdk) — supported platforms, GPUs,
|
||||
CloudXR/OpenXR runtime versions, and headsets.
|
||||
@@ -6,11 +6,12 @@ Encoding frames into an MP4 is a full FFmpeg pipeline: choice of encoder, pixel
|
||||
|
||||
You can set these parameters from the CLI with `--dataset.rgb_encoder.<field>` (e.g. with `lerobot-record` or `lerobot-rollout`). The same block applies to every camera video stream in that run.
|
||||
|
||||
> [!TIP]
|
||||
> Video storage must be on for `rgb_encoder` to have any effect —
|
||||
> `use_videos=True` in Python APIs, or `--dataset.video=true` on the CLI (the
|
||||
> recording default). With video off, inputs stay as images and `rgb_encoder` is
|
||||
> ignored.
|
||||
<Tip>
|
||||
Video storage must be on for `rgb_encoder` to have any effect —
|
||||
`use_videos=True` in Python APIs, or `--dataset.video=true` on the CLI (the
|
||||
recording default). With video off, inputs stay as images and `rgb_encoder` is
|
||||
ignored.
|
||||
</Tip>
|
||||
|
||||
For details on **when** frames are written vs. encoded (streaming vs. post-episode), queues, and other top-level `--dataset.*` switches, see [Streaming Video Encoding](./streaming_video_encoding). For an encoding-parameter comparison and experiments, see the [video-benchmark Space](https://huggingface.co/spaces/lerobot/video-benchmark).
|
||||
|
||||
@@ -42,10 +43,12 @@ lerobot-record \
|
||||
|
||||
## Tuning parameters
|
||||
|
||||
> [!WARNING]
|
||||
> The defaults are tuned to balance **compression ratio**, **visual quality**, and **decoding/seek speed** for typical robotics datasets. Changing them can affect both recording (CPU load, frame drops) and training (decoding throughput, image quality).
|
||||
>
|
||||
> Only override these parameters if you have a specific reason to, and measure the impact on your pipeline before relying on the new settings.
|
||||
<Tip warning={true}>
|
||||
The defaults are tuned to balance **compression ratio**, **visual quality**, and **decoding/seek speed** for typical robotics datasets. Changing them can affect both recording (CPU load, frame drops) and training (decoding throughput, image quality).
|
||||
|
||||
Only override these parameters if you have a specific reason to, and measure the impact on your pipeline before relying on the new settings.
|
||||
|
||||
</Tip>
|
||||
|
||||
All flags below are prefixed with `--dataset.rgb_encoder.` on the CLI.
|
||||
|
||||
@@ -66,92 +69,25 @@ All flags below are prefixed with `--dataset.rgb_encoder.` on the CLI.
|
||||
|
||||
Depth maps (Intel RealSense, Reachy 2) are stored as their **own video streams** alongside the RGB streams. Raw depth (`uint16` millimetres or `float32` metres) can't survive an 8-bit codec, so LeRobot **quantizes** each map to a 12-bit code (`[0, 4095]`) — logarithmically by default, to match the `1/depth` error profile of depth sensors — then packs it into a high-bit-depth pixel format (`gray12le`) and encodes it with a 12-bit codec.
|
||||
|
||||
<div style="margin:28px 0;padding:14px 0;">
|
||||
<div style="margin:0 auto;display:flex;flex-wrap:wrap;justify-content:center;align-items:stretch;gap:6px;font-family:'Source Sans 3',ui-sans-serif,system-ui,sans-serif;font-size:14px;font-weight:600;color:#1B1B1D;">
|
||||
<span style="display:flex;flex-direction:column;justify-content:center;align-items:center;text-align:center;gap:2px;background:#DBEAFE;color:#1D4ED8;border-radius:9px;padding:8px 12px;">
|
||||
<span>Raw depth</span>
|
||||
<span style="font-size:11px;font-weight:400;color:#3B6FD4;white-space:nowrap;">
|
||||
uint16 mm
|
||||
<br />
|
||||
float32 m
|
||||
</span>
|
||||
</span>
|
||||
<span style="display:flex;align-items:center;font-size:16px;color:#C3CBD9;">
|
||||
→
|
||||
</span>
|
||||
<div style="border:2px dashed #C4B5FD;border-radius:13px;padding:18px 12px 12px;position:relative;display:flex;align-items:stretch;gap:6px;">
|
||||
<span style="position:absolute;top:-10px;left:12px;background:#fff;padding:0 6px;font-size:11px;font-weight:700;color:#7E22CE;text-transform:uppercase;letter-spacing:0.5px;white-space:nowrap;">
|
||||
Record time
|
||||
</span>
|
||||
<span style="display:flex;flex-direction:column;justify-content:center;align-items:center;text-align:center;gap:2px;background:#F3E8FF;color:#7E22CE;border-radius:9px;padding:8px 12px;">
|
||||
<span>Clip</span>
|
||||
<span style="font-size:11px;font-weight:400;color:#9061C2;white-space:nowrap;">
|
||||
to [depth_min,
|
||||
<br />
|
||||
depth_max]
|
||||
</span>
|
||||
</span>
|
||||
<span style="display:flex;align-items:center;font-size:16px;color:#C3CBD9;">
|
||||
→
|
||||
</span>
|
||||
<span style="display:flex;flex-direction:column;justify-content:center;align-items:center;text-align:center;gap:2px;background:#F3E8FF;color:#7E22CE;border-radius:9px;padding:8px 12px;">
|
||||
<span>Quantize</span>
|
||||
<span style="font-size:11px;font-weight:400;color:#9061C2;white-space:nowrap;">
|
||||
12-bit codes 0–4095
|
||||
<br />
|
||||
log (default) or linear
|
||||
</span>
|
||||
</span>
|
||||
<span style="display:flex;align-items:center;font-size:16px;color:#C3CBD9;">
|
||||
→
|
||||
</span>
|
||||
<span style="display:flex;flex-direction:column;justify-content:center;align-items:center;text-align:center;gap:2px;background:#F3E8FF;color:#7E22CE;border-radius:9px;padding:8px 12px;">
|
||||
<span>Pack</span>
|
||||
<span style="font-size:11px;font-weight:400;color:#9061C2;white-space:nowrap;">
|
||||
into gray12le
|
||||
<br />
|
||||
plane
|
||||
</span>
|
||||
</span>
|
||||
<span style="display:flex;align-items:center;font-size:16px;color:#C3CBD9;">
|
||||
→
|
||||
</span>
|
||||
<span style="display:flex;flex-direction:column;justify-content:center;align-items:center;text-align:center;gap:2px;background:#F3E8FF;color:#7E22CE;border-radius:9px;padding:8px 12px;">
|
||||
<span>Encode</span>
|
||||
<span style="font-size:11px;font-weight:400;color:#9061C2;white-space:nowrap;">
|
||||
HEVC
|
||||
<br />
|
||||
Main 12
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
<span style="display:flex;align-items:center;font-size:16px;color:#C3CBD9;">
|
||||
→
|
||||
</span>
|
||||
<span style="display:flex;flex-direction:column;justify-content:center;align-items:center;text-align:center;gap:2px;background:#FEF3C7;color:#B45309;border-radius:9px;padding:8px 12px;">
|
||||
<span>MP4</span>
|
||||
<span style="font-size:11px;font-weight:400;color:#C77D18;white-space:nowrap;">
|
||||
stored
|
||||
<br />
|
||||
stream
|
||||
</span>
|
||||
</span>
|
||||
<span style="display:flex;align-items:center;font-size:16px;color:#34A06B;">
|
||||
→
|
||||
</span>
|
||||
<div style="border:2px dashed #6EE7B7;border-radius:13px;padding:18px 12px 12px;position:relative;display:flex;align-items:center;gap:6px;">
|
||||
<span style="position:absolute;top:-10px;left:12px;background:#fff;padding:0 6px;font-size:11px;font-weight:700;color:#047857;text-transform:uppercase;letter-spacing:0.5px;white-space:nowrap;">
|
||||
Load time
|
||||
</span>
|
||||
<span style="display:flex;flex-direction:column;justify-content:center;align-items:center;text-align:center;gap:2px;background:#D1FAE5;color:#047857;border-radius:9px;padding:8px 12px;">
|
||||
<span>Dequantize</span>
|
||||
<span style="font-size:11px;font-weight:400;color:#059669;white-space:nowrap;">
|
||||
to mm / m
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A["Raw depth (uint16 mm / float32 m)"] --> B["Clip to depth_min, depth_max"]
|
||||
B --> C["Quantize to 12-bit code 0–4095 (log or linear)"]
|
||||
C --> D["Pack into gray12le"]
|
||||
D --> E["Encode video (hevc Main 12)"]
|
||||
E --> F[("MP4 + metadata: depth_min/max, shift, use_log")]
|
||||
F -. "load time (depth_output_unit)" .-> G["Dequantize to mm or m"]
|
||||
|
||||
classDef input fill:#e3f2fd,stroke:#1565c0,color:#0d47a1;
|
||||
classDef encode fill:#ede7f6,stroke:#5e35b1,color:#311b92;
|
||||
classDef store fill:#fff8e1,stroke:#f9a825,color:#e65100;
|
||||
classDef load fill:#e8f5e9,stroke:#2e7d32,color:#1b5e20;
|
||||
|
||||
class A input;
|
||||
class B,C,D,E encode;
|
||||
class F store;
|
||||
class G load;
|
||||
```
|
||||
|
||||
Configure the depth pipeline through a parallel **`depth_encoder`** block (`DepthEncoderConfig`). It shares every `RGBEncoderConfig` field (`vcodec`, `pix_fmt`, `crf`, …) and adds four quantizer knobs, set via `--dataset.depth_encoder.<field>`:
|
||||
|
||||
@@ -232,16 +168,15 @@ After the first episode of a video stream is encoded, the encoder configuration
|
||||
|
||||
Two sources contribute to the `info` block:
|
||||
|
||||
| Source | Where it comes from | Fields |
|
||||
| ------------------- | ----------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------- |
|
||||
| **Stream-derived** | Read back from the encoded MP4 with PyAV. | `video.height`, `video.width`, `video.codec`, `video.pix_fmt`, `video.fps`, `video.channels`, `is_depth_map`, `audio.*` |
|
||||
| **Encoder-derived** | Taken from `RGBEncoderConfig` / `DepthEncoderConfig`. | `video.g`, `video.crf`, `video.preset`, `video.fast_decode`, `video.video_backend`, `video.extra_options` |
|
||||
- **Stream-derived** (read back from the encoded MP4 with PyAV): `video.height`, `video.width`, `video.codec`, `video.pix_fmt`, `video.fps`, `video.channels`, `is_depth_map`, plus `audio.*` if an audio stream is present.
|
||||
- **Encoder-derived** (taken from `RGBEncoderConfig` or `DepthEncoderConfig`): `video.g`, `video.crf`, `video.preset`, `video.fast_decode`, `video.video_backend`, `video.extra_options`.
|
||||
|
||||
> [!IMPORTANT]
|
||||
> This block is populated **once**, from the **first** episode. It assumes every
|
||||
> episode in the dataset was encoded with the same `rgb_encoder`. Changing
|
||||
> encoder settings partway through a recording is not supported — the
|
||||
> `info.json` will only reflect the parameters used for the first episode.
|
||||
<Tip>
|
||||
This block is populated **once**, from the **first** episode. It assumes every
|
||||
episode in the dataset was encoded with the same `rgb_encoder`. Changing
|
||||
encoder settings partway through a recording is not supported — the
|
||||
`info.json` will only reflect the parameters used for the first episode.
|
||||
</Tip>
|
||||
|
||||
---
|
||||
|
||||
@@ -249,7 +184,5 @@ Two sources contribute to the `info` block:
|
||||
|
||||
When aggregating datasets with `merge_datasets`, video files are concatenated as-is (no re-encoding), and encoder fields in `info.json` are merged per-key:
|
||||
|
||||
| Merge rule | Fields | Behaviour |
|
||||
| ------------------ | ---------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| **Must match** | `video.codec`, `video.pix_fmt`, `video.height`, `video.width`, `video.fps` | Stream-derived fields must match across sources, otherwise FFmpeg's concat demuxer fails. |
|
||||
| **Merged loosely** | `video.g`, `video.crf`, `video.preset`, `video.fast_decode`, `video.extra_options` | Encoder-tuning fields. If every source agrees, the value is kept; if not, it's set to `null` (or `{}` for `video.extra_options`) and a warning is logged. |
|
||||
- **Stream-derived fields must match** across sources: `video.codec`, `video.pix_fmt`, `video.height`, `video.width`, `video.fps`. Otherwise FFmpeg's concat demuxer fails.
|
||||
- **Encoder-tuning fields are merged loosely**: `video.g`, `video.crf`, `video.preset`, `video.fast_decode`, `video.extra_options`. If every source agrees, the value is kept; if not, it's set to `null` (or `{}` for `video.extra_options`) and a warning is logged.
|
||||
|
||||
@@ -46,11 +46,8 @@ CMD = (
|
||||
"apt-get update -qq && apt-get install -y -qq git ffmpeg && "
|
||||
"pip install --no-deps "
|
||||
"'lerobot @ git+https://github.com/huggingface/lerobot.git@main' && "
|
||||
# Pins mirror pyproject.toml — unpinned installs pull av 18 / datasets 5 /
|
||||
# draccus 0.11, which break lerobot at import time.
|
||||
"pip install --upgrade-strategy only-if-needed "
|
||||
"'datasets>=4.7.0,<5.0.0' 'pyarrow>=21.0.0,<30.0.0' 'av>=15.0.0,<16.0.0' 'draccus==0.10.0' "
|
||||
"'pandas>=2.0.0,<3.0.0' jsonlines gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
|
||||
"datasets pyarrow av jsonlines draccus gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
|
||||
"openai && "
|
||||
"export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
|
||||
"export VLLM_VIDEO_BACKEND=pyav && "
|
||||
|
||||
@@ -1,131 +0,0 @@
|
||||
# Isaac Teleop → SO-101
|
||||
|
||||
Teleoperate an SO-101/SO-100 follower arm — and record LeRobot datasets — with NVIDIA
|
||||
[Isaac Teleop](https://github.com/NVIDIA/IsaacTeleop). Two input devices ship today:
|
||||
|
||||
- **XR (VR) controller** (`--teleop.type=xr_controller`) — the controller's grip pose drives the
|
||||
end-effector through a squeeze-to-engage clutch and LeRobot's Cartesian IK pipeline; the analog
|
||||
trigger drives the gripper.
|
||||
- **SO-101 leader arm** (`--teleop.type=so101_leader`) — a back-drivable leader arm mirrored 1:1
|
||||
onto the follower via Isaac Teleop's native `so101_leader` plugin (no clutch, no IK).
|
||||
|
||||
The full narrative guide (how the clutch works, CloudXR setup, headset pairing, tuning, and
|
||||
troubleshooting) is in the [LeRobot docs](https://huggingface.co/docs/lerobot/isaac_teleop)
|
||||
(source: `docs/source/isaac_teleop.mdx`). This README is the canonical install and usage
|
||||
reference.
|
||||
|
||||
## Requirements
|
||||
|
||||
- Linux workstation (see NVIDIA's
|
||||
[system requirements](https://nvidia.github.io/IsaacTeleop/main/references/requirements.html)
|
||||
for supported OS/GPU/headset combinations; `isaacteleop` publishes Linux wheels only).
|
||||
- An SO-101 (or SO-100) follower arm, calibrated with `lerobot-calibrate`.
|
||||
- For the XR device: a CloudXR-capable headset (e.g. Quest 3, Pico 4, Apple Vision Pro) on the
|
||||
same network.
|
||||
- For the leader device: a second, back-drivable SO-101 leader arm and the `so101_leader` plugin
|
||||
binary built from the Isaac Teleop source tree (see
|
||||
[Build from source](https://nvidia.github.io/IsaacTeleop/main/getting_started/build_from_source/index.html)).
|
||||
|
||||
## Installation
|
||||
|
||||
This example lives in the LeRobot repository and is not part of the `lerobot` pip package, so
|
||||
work from a source checkout. From the repo root:
|
||||
|
||||
```bash
|
||||
# LeRobot with the extras this example uses:
|
||||
# feetech - SO-101 serial motor bus
|
||||
# kinematics - Placo IK solver (XR controller path)
|
||||
# dataset - dataset recording (record.py)
|
||||
# huggingface_hub >= 1.5 is needed by the automatic URDF fetch (Buckets API).
|
||||
uv pip install -e ".[feetech,kinematics,dataset]" "huggingface_hub>=1.5"
|
||||
|
||||
# Isaac Teleop from public PyPI. `cloudxr` brings the CloudXR runtime bindings;
|
||||
# `retargeters-lite` is the scipy-based retargeter path that resolves on both
|
||||
# x86_64 and ARM (the full `retargeters` extra does not resolve on aarch64).
|
||||
uv pip install "isaacteleop[cloudxr,retargeters-lite]~=1.3.131" "scipy>=1.14"
|
||||
|
||||
# Optional, x86_64 only: the full retargeter stack.
|
||||
uv pip install "isaacteleop[retargeters]~=1.3.131"
|
||||
```
|
||||
|
||||
One-time CloudXR EULA (the auto-launch prompts on stdin and would hang on a headless machine):
|
||||
|
||||
```bash
|
||||
python -m isaacteleop.cloudxr --accept-eula
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
Run everything from the repo root with `python -m` so the `examples` package resolves.
|
||||
|
||||
### Teleoperate — XR controller
|
||||
|
||||
```bash
|
||||
python -m examples.isaac_teleop_to_so101.teleoperate \
|
||||
--robot.type=so101_follower \
|
||||
--robot.port=/dev/ttyACM0 \
|
||||
--robot.id=so101_follower_arm \
|
||||
--teleop.type=xr_controller
|
||||
```
|
||||
|
||||
On startup the script launches the CloudXR runtime (~30 s), prints the workstation IP to enter in
|
||||
the headset's CloudXR web client, waits for the controllers to stream, slews the arm to a reset
|
||||
pose (`--reset_to_origin=false` to skip), and then: **hold the squeeze/grip** to engage, move the
|
||||
controller to drive the arm, pull the trigger to close the gripper. Releasing the squeeze freezes
|
||||
the arm. The SO-101 URDF is fetched automatically from the `lerobot/robot-urdfs` Hugging Face
|
||||
bucket into the LeRobot cache on first run.
|
||||
|
||||
To customize the reset pose: back-drive the arm to the pose you want, then
|
||||
|
||||
```bash
|
||||
python -m examples.isaac_teleop_to_so101.override_reset_pose --port /dev/ttyACM0 --id so101_follower_arm
|
||||
```
|
||||
|
||||
which writes it to `HF_LEROBOT_HOME/reset_poses/<robot.name>/<robot.id>.json`; runs with the same
|
||||
`--robot.id` use it automatically.
|
||||
|
||||
### Teleoperate — SO-101 leader arm
|
||||
|
||||
```bash
|
||||
python -m examples.isaac_teleop_to_so101.teleoperate \
|
||||
--robot.type=so101_follower --robot.port=/dev/ttyACM0 --robot.id=so101_follower_arm \
|
||||
--teleop.type=so101_leader --teleop.port=/dev/ttyACM1 --teleop.id=so101_leader_arm \
|
||||
--launch_plugin=/path/to/IsaacTeleop/install/plugins/so101_leader/so101_leader_plugin
|
||||
```
|
||||
|
||||
The follower is first slewed to the leader's pose over `--align_duration` seconds
|
||||
(`--align=false` to skip), then mirrors it 1:1. The plugin reuses the serial leader's calibration
|
||||
(`HF_LEROBOT_CALIBRATION/teleoperators/so_leader/<teleop.id>.json`).
|
||||
|
||||
### Record a dataset
|
||||
|
||||
`record.py` takes the same `--robot.*`/`--teleop.*`/loop flags plus `lerobot-record`-style
|
||||
`--dataset.*` flags:
|
||||
|
||||
```bash
|
||||
python -m examples.isaac_teleop_to_so101.record \
|
||||
--robot.type=so101_follower --robot.port=/dev/ttyACM0 --robot.id=so101_follower_arm \
|
||||
--teleop.type=xr_controller \
|
||||
--robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
|
||||
--dataset.repo_id=<hf_user>/<dataset_name> \
|
||||
--dataset.single_task="Pick up the cube" \
|
||||
--dataset.num_episodes=3 --dataset.episode_time_s=20 --dataset.reset_time_s=5
|
||||
```
|
||||
|
||||
Keyboard shortcuts (terminal-first, so they work over SSH): **Right/n** end episode early,
|
||||
**Left/r** re-record, **Esc/q** stop after the current episode.
|
||||
|
||||
Run either script with `--help` for all flags.
|
||||
|
||||
## Layout
|
||||
|
||||
```
|
||||
isaac_teleop/ device library: session lifecycle (base.py), XRController,
|
||||
SO101LeaderArm, Clutch, configs, and the XR→IK processor step
|
||||
common.py shared loop infra: device bundles, clutch/IK pipeline wiring,
|
||||
reset/align slews, URDF fetch, keyboard listener
|
||||
teleoperate.py teleoperation CLI (device selected via --teleop.type)
|
||||
record.py dataset-recording CLI (same device selection + --dataset.*)
|
||||
override_reset_pose.py save the current joints as the per-arm reset pose
|
||||
default.env CloudXR device-profile overrides passed to the launcher
|
||||
```
|
||||
@@ -1,17 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 NVIDIA Corporation and The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Isaac Teleop -> SO-101 example package."""
|
||||
@@ -1,650 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 NVIDIA Corporation and The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Shared device + control-loop infrastructure for the Isaac Teleop -> SO-101 examples.
|
||||
|
||||
Consumed by ``teleoperate.py`` and ``record.py``, which both build a per-device
|
||||
:class:`Device` bundle and run the same loop: read -> (maybe command) -> hold-when-idle ->
|
||||
sleep. A :class:`Device` bundles three closures: ``compute(obs) -> RobotAction | None``
|
||||
(``None`` = hold at the measured pose while idle), ``startup``, and ``cleanup``. The devices:
|
||||
|
||||
* ``xr_controller`` — a thin :class:`XRController` whose raw grip pose an in-loop
|
||||
:class:`Clutch` turns into an EE target for LeRobot's Cartesian IK pipeline.
|
||||
* ``so101_leader`` — a back-drivable leader arm mirrored 1:1 into the follower.
|
||||
|
||||
Requires the ``isaacteleop`` package and an OpenXR runtime (install instructions in this
|
||||
folder's ``README.md``). User-facing guide: ``docs/source/isaac_teleop.mdx``.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from importlib.resources import files
|
||||
from pathlib import Path
|
||||
from typing import Protocol
|
||||
|
||||
import numpy as np
|
||||
|
||||
from lerobot.model.kinematics import RobotKinematics
|
||||
from lerobot.processor import (
|
||||
RobotProcessorPipeline,
|
||||
robot_action_observation_to_transition,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
from lerobot.robots import RobotConfig, make_robot_from_config
|
||||
from lerobot.robots.so_follower import SOFollowerConfig # noqa: F401 (registers so101_follower)
|
||||
from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||
EEBoundsAndSafety,
|
||||
InverseKinematicsEEToJoints,
|
||||
)
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, HF_LEROBOT_HOME, TELEOPERATORS
|
||||
from lerobot.utils.robot_utils import precise_sleep
|
||||
|
||||
from .isaac_teleop import (
|
||||
Clutch,
|
||||
IsaacTeleopConfig,
|
||||
MapXRControllerActionToRobotAction,
|
||||
SO101LeaderArm,
|
||||
SO101LeaderArmConfig,
|
||||
XRController,
|
||||
)
|
||||
|
||||
# Fixed rate [Hz] for the teleoperate loop and the pre-loop slews / connect-wait poll sleeps.
|
||||
FPS = 30
|
||||
|
||||
# CloudXR device-profile env file passed to the launcher (see default.env in this package).
|
||||
CLOUDXR_ENV_FILE = str(files(__package__) / "default.env")
|
||||
|
||||
|
||||
class LoopConfig(Protocol):
|
||||
"""Structural type for the loop/launch knobs ``build_device`` and the ``setup_*`` read.
|
||||
|
||||
Both ``TeleoperateConfig`` and ``RecordConfig`` satisfy it, keeping ``common`` decoupled
|
||||
from either entry point's concrete config.
|
||||
"""
|
||||
|
||||
teleop: IsaacTeleopConfig
|
||||
robot: RobotConfig
|
||||
launch_plugin: str | None
|
||||
reset_to_origin: bool
|
||||
reset_duration: float
|
||||
align: bool
|
||||
align_duration: float
|
||||
|
||||
|
||||
# Per-device bundle consumed by the shared loop. ``compute`` returns None to mean
|
||||
# "idle -> hold at the measured pose"; ``startup`` warms up; ``cleanup`` reaps/disconnects.
|
||||
@dataclass(frozen=True)
|
||||
class Device:
|
||||
compute: Callable[[RobotObservation | None], RobotAction | None]
|
||||
startup: Callable[[], None]
|
||||
cleanup: Callable[[], None]
|
||||
|
||||
|
||||
def hold_action(obs: RobotObservation, motor_names: list[str]) -> dict[str, float]:
|
||||
"""Re-send the measured joints — the explicit hold when a device is idle."""
|
||||
return {f"{name}.pos": float(obs[f"{name}.pos"]) for name in motor_names}
|
||||
|
||||
|
||||
class HoldLatch:
|
||||
"""Resolve the per-frame action, holding one LATCHED pose while the device is idle.
|
||||
|
||||
Re-sending the freshly measured joints on every idle frame would ratchet the arm
|
||||
downward: under gravity the P-only servo settles below its goal by a steady-state
|
||||
error, so each re-command of the measurement lowers the goal by that error again.
|
||||
Latching the target once on the active->idle transition holds a fixed pose instead.
|
||||
"""
|
||||
|
||||
def __init__(self, motor_names: list[str]):
|
||||
self._motor_names = motor_names
|
||||
self._held: dict[str, float] | None = None
|
||||
|
||||
def resolve(self, action: RobotAction | None, obs: RobotObservation) -> RobotAction:
|
||||
"""Pass through an active action (clearing the latch); latch + hold when idle."""
|
||||
if action is not None:
|
||||
self._held = None
|
||||
return action
|
||||
if self._held is None:
|
||||
self._held = hold_action(obs, self._motor_names)
|
||||
return self._held
|
||||
|
||||
|
||||
def slew(
|
||||
robot,
|
||||
motor_names: list[str],
|
||||
target_fn: Callable[[], dict[str, float]],
|
||||
duration_s: float,
|
||||
) -> None:
|
||||
"""Linearly slew all joints from their current measured pose toward a target.
|
||||
|
||||
``target_fn`` is called EACH step, so the leader can pass a live re-read (landing on its
|
||||
current pose at ``alpha == 1`` for a continuous handoff) while XR passes a constant.
|
||||
"""
|
||||
obs = robot.get_observation()
|
||||
start = {name: float(obs[f"{name}.pos"]) for name in motor_names}
|
||||
n_steps = max(1, int(duration_s * FPS))
|
||||
for step in range(1, n_steps + 1):
|
||||
alpha = step / n_steps
|
||||
target = target_fn()
|
||||
action = {f"{name}.pos": start[name] + alpha * (target[name] - start[name]) for name in motor_names}
|
||||
robot.send_action(action)
|
||||
precise_sleep(1.0 / FPS)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# XR controller device
|
||||
# ============================================================================
|
||||
|
||||
# Per-frame EE rate limit [m]. With raise_on_jump=False, EEBoundsAndSafety clamps an
|
||||
# over-limit step instead of raising, absorbing a tracking glitch as one slow frame. At
|
||||
# FPS=30, 0.1 m/frame caps EE speed at ~3 m/s. (end_effector_bounds clips the absolute target.)
|
||||
MAX_EE_STEP_M = 0.1
|
||||
|
||||
# Soft-orientation IK weight: small but nonzero so the wrist follows the hand while position
|
||||
# dominates (the 5-DOF SO-101 cannot realize an arbitrary orientation). 0.0 = position-only.
|
||||
IK_ORIENTATION_WEIGHT = 0.01
|
||||
|
||||
|
||||
def _ensure_so101_urdf() -> str:
|
||||
"""Return the cached SO-101 URDF path, fetching the ``so101`` folder (URDF + meshes) from
|
||||
the public ``lerobot/robot-urdfs`` HF bucket into the LeRobot cache on first use."""
|
||||
dest_dir = HF_LEROBOT_HOME / "robot-urdfs" / "so101"
|
||||
urdf_path = dest_dir / "so101_new_calib.urdf"
|
||||
# Completeness marker written only after a FULL sync: the URDF file alone is not a
|
||||
# completeness signal (an interrupted first sync can leave the meshes it references
|
||||
# missing, which the URDF's mere existence would then hide forever). Re-syncing is
|
||||
# idempotent and repairs a partial cache; delete the folder to force a re-download.
|
||||
marker = dest_dir / ".sync_complete"
|
||||
if not marker.exists():
|
||||
from huggingface_hub import sync_bucket
|
||||
|
||||
sync_bucket("hf://buckets/lerobot/robot-urdfs/so101", str(dest_dir), quiet=True)
|
||||
marker.touch()
|
||||
return str(urdf_path)
|
||||
|
||||
|
||||
# Default duration [s] for the startup reset-to-origin slew.
|
||||
RESET_DURATION_S = 5.0
|
||||
|
||||
# Optional cached file written by override_reset_pose.py. When present it takes priority over RESET_ORIGIN_DEG.
|
||||
RESET_POSE_FILE = str(HF_LEROBOT_HOME / "reset_poses" / "{robot_name}" / "{robot_id}.json")
|
||||
|
||||
# Reset target in each motor's native units (arm joints in degrees, gripper RANGE_0_100,
|
||||
# 100 = open). An empirically comfortable pose (elbow/wrist bent) avoiding the singularity of
|
||||
# a fully-extended arm; assumes standard calibration. Override per-arm via override_reset_pose.py.
|
||||
RESET_ORIGIN_DEG: dict[str, float] = {
|
||||
"shoulder_pan": -4.0,
|
||||
"shoulder_lift": -103.0,
|
||||
"elbow_flex": 97.0,
|
||||
"wrist_flex": 78.0,
|
||||
"wrist_roll": -65.0,
|
||||
"gripper": 0.0,
|
||||
}
|
||||
|
||||
|
||||
def _load_reset_target(reset_pose_file: Path, motor_names: list[str]) -> dict[str, float]:
|
||||
"""Return reset targets: the saved reset pose if present, else RESET_ORIGIN_DEG."""
|
||||
if reset_pose_file.exists():
|
||||
saved = json.loads(reset_pose_file.read_text())
|
||||
# Fill any missing motors from the fallback dict.
|
||||
return {name: float(saved.get(name, RESET_ORIGIN_DEG.get(name, 0.0))) for name in motor_names}
|
||||
return {name: RESET_ORIGIN_DEG.get(name, 0.0) for name in motor_names}
|
||||
|
||||
|
||||
# CloudXR web client URL opened in the headset (Isaac Teleop quick start, step 5).
|
||||
_CLOUDXR_WEB_CLIENT_URL = "https://nvidia.github.io/IsaacTeleop/client"
|
||||
# WSS-proxy / self-signed-cert port the operator accepts in-browser before connecting.
|
||||
_CLOUDXR_WSS_PORT = 48322
|
||||
# How often to re-print the connection hint while waiting for the headset [s].
|
||||
_XR_CONNECT_REMINDER_S = 15.0
|
||||
# Virtual / bridge / USB-gadget interfaces a headset can't reach over the network — skip
|
||||
# by name prefix (``docker0``, compose ``br-*``, ``veth*``, libvirt ``virbr*``, and the
|
||||
# Tegra USB device-mode bridge ``l4tbr0``).
|
||||
_SKIP_IFACE_PREFIXES = ("docker", "br-", "veth", "virbr", "l4tbr")
|
||||
|
||||
|
||||
def _primary_ipv4() -> str | None:
|
||||
"""The workstation's primary outbound IPv4, via the UDP-socket trick (``connect()`` on a
|
||||
datagram socket selects the egress interface without sending packets)."""
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
|
||||
try:
|
||||
s.connect(("8.8.8.8", 80))
|
||||
return s.getsockname()[0]
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
|
||||
def _candidate_ipv4s() -> list[tuple[str, str]]:
|
||||
"""Return ``[(interface, ipv4), ...]`` the headset might reach this workstation at.
|
||||
|
||||
Lists each interface's IPv4 via ``psutil`` (dropping loopback, link-local, and the
|
||||
virtual/bridge interfaces in ``_SKIP_IFACE_PREFIXES``), primary outbound first. Falls
|
||||
back to just the primary IP when ``psutil`` is unavailable.
|
||||
"""
|
||||
primary = _primary_ipv4()
|
||||
found: list[tuple[str, str]] = []
|
||||
try:
|
||||
import psutil
|
||||
|
||||
for iface, addrs in psutil.net_if_addrs().items():
|
||||
if iface.startswith(_SKIP_IFACE_PREFIXES):
|
||||
continue
|
||||
for addr in addrs:
|
||||
if addr.family != socket.AF_INET:
|
||||
continue
|
||||
ip = addr.address
|
||||
if ip.startswith("127.") or ip.startswith("169.254."):
|
||||
continue
|
||||
found.append((iface, ip))
|
||||
except Exception:
|
||||
if primary:
|
||||
found.append(("default", primary))
|
||||
found.sort(key=lambda t: t[1] != primary) # primary outbound interface first
|
||||
return found
|
||||
|
||||
|
||||
def _print_xr_connect_help() -> None:
|
||||
"""Print how to connect the headset to this workstation over CloudXR."""
|
||||
ips = _candidate_ipv4s()
|
||||
print("\n" + "=" * 76)
|
||||
print("Connect your XR headset to this workstation over NVIDIA CloudXR:")
|
||||
print(f" 1. In the headset, open the CloudXR web client: {_CLOUDXR_WEB_CLIENT_URL}")
|
||||
print(" 2. Enter this workstation's IP address:")
|
||||
if ips:
|
||||
for iface, ip in ips:
|
||||
print(f" {ip:<15} ({iface})")
|
||||
if len(ips) > 1:
|
||||
print(" (use the address on the same network as your headset)")
|
||||
else:
|
||||
print(" <could not determine — check `hostname -I` / `ip addr`>")
|
||||
print(f" 3. Accept the self-signed cert at https://<that-ip>:{_CLOUDXR_WSS_PORT}/ , then Connect.")
|
||||
print("=" * 76 + "\n")
|
||||
|
||||
|
||||
def _wait_for_xr_controller(teleop_device: XRController) -> None:
|
||||
"""Block until the XR controller is tracked, polling ``get_action()`` and re-printing a
|
||||
reminder every ``_XR_CONNECT_REMINDER_S``. User-paced; ``Ctrl-C`` aborts (no hard timeout).
|
||||
"""
|
||||
_print_xr_connect_help()
|
||||
print("Waiting for the headset controllers to start streaming… (Ctrl-C to abort)")
|
||||
last_reminder = time.time()
|
||||
while True:
|
||||
teleop_device.get_action() # steps the session; updates is_tracking
|
||||
if teleop_device.is_tracking:
|
||||
print("Headset connected — controllers are streaming.")
|
||||
return
|
||||
if time.time() - last_reminder >= _XR_CONNECT_REMINDER_S:
|
||||
print("…still waiting for the headset to connect (Ctrl-C to abort).")
|
||||
last_reminder = time.time()
|
||||
time.sleep(1.0 / FPS)
|
||||
|
||||
|
||||
def setup_xr(cfg: LoopConfig, robot, motor_names: list[str]) -> Device:
|
||||
"""Build the XR controller device bundle (clutch + soft-orientation IK pipeline)."""
|
||||
kinematics_solver = RobotKinematics(
|
||||
urdf_path=_ensure_so101_urdf(),
|
||||
target_frame_name="gripper_frame_link",
|
||||
joint_names=motor_names,
|
||||
)
|
||||
|
||||
teleop_config = cfg.teleop # XRControllerConfig (selected via --teleop.type=xr_controller)
|
||||
teleop_device = XRController(teleop_config)
|
||||
|
||||
# The clutch (below) turns the raw grip pose into an absolute base-frame ee_pose; this
|
||||
# pipeline maps it to joint targets: rename -> bounds/rate-limit -> IK.
|
||||
xr_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
|
||||
steps=[
|
||||
MapXRControllerActionToRobotAction(),
|
||||
# raise_on_jump=False: an over-limit step (e.g. a tracking glitch) is clamped +
|
||||
# warned instead of raised, since a crash mid-loop would leave the arm uncontrolled.
|
||||
# z floor 0.0 keeps a stray target above the table; x/y stay at a loose [-1,1]m box.
|
||||
EEBoundsAndSafety(
|
||||
end_effector_bounds={"min": [-1.0, -1.0, 0.0], "max": [1.0, 1.0, 1.0]},
|
||||
max_ee_step_m=MAX_EE_STEP_M,
|
||||
raise_on_jump=False,
|
||||
),
|
||||
# initial_guess_current_joints=False: warm-start from the previous IK solution so
|
||||
# the joint trajectory stays continuous frame-to-frame.
|
||||
InverseKinematicsEEToJoints(
|
||||
kinematics=kinematics_solver,
|
||||
motor_names=motor_names,
|
||||
initial_guess_current_joints=False,
|
||||
orientation_weight=IK_ORIENTATION_WEIGHT,
|
||||
),
|
||||
],
|
||||
to_transition=robot_action_observation_to_transition,
|
||||
to_output=transition_to_robot_action,
|
||||
)
|
||||
|
||||
# The clutch is built in startup() (after the optional reset slew, seeded from the
|
||||
# post-slew MEASURED pose) and shared with compute() via nonlocal.
|
||||
clutch: Clutch | None = None
|
||||
prev_enabled = False
|
||||
|
||||
def startup() -> None:
|
||||
nonlocal clutch
|
||||
# Connect and wait for the operator to don the headset BEFORE moving the arm, so the
|
||||
# reset slew happens while they are watching in VR.
|
||||
teleop_device.connect()
|
||||
if not teleop_device.is_connected:
|
||||
raise ValueError("Teleop is not connected!")
|
||||
_wait_for_xr_controller(teleop_device)
|
||||
|
||||
if cfg.reset_to_origin:
|
||||
reset_pose_file = Path(RESET_POSE_FILE.format(robot_name=robot.name, robot_id=robot.id))
|
||||
target = _load_reset_target(reset_pose_file, motor_names)
|
||||
source = str(reset_pose_file) if reset_pose_file.exists() else "hardcoded defaults"
|
||||
print(f"Reset target source: {source}")
|
||||
print(f"Resetting to origin over {cfg.reset_duration:.1f} s…")
|
||||
slew(robot, motor_names, lambda: target, cfg.reset_duration)
|
||||
print("Reset complete.")
|
||||
|
||||
# Seed the clutch home from the arm's measured pose (FK of the current joints) so the
|
||||
# first engage is jump-free, whether or not a reset slew ran.
|
||||
obs0 = robot.get_observation()
|
||||
q_measured_deg = np.array([float(obs0[f"{name}.pos"]) for name in motor_names], dtype=float)
|
||||
home_base_T_ee = kinematics_solver.forward_kinematics(q_measured_deg) # noqa: N806
|
||||
clutch = Clutch(home_base_T_ee)
|
||||
|
||||
print("Starting teleop loop. Squeeze and move the controller to teleoperate the robot...")
|
||||
|
||||
def compute(robot_obs: RobotObservation | None) -> RobotAction | None:
|
||||
nonlocal prev_enabled
|
||||
if clutch is None: # set in startup(), which runs before compute()
|
||||
raise RuntimeError("compute() called before startup(); the clutch is not initialized")
|
||||
xr_action = teleop_device.get_action()
|
||||
grip_pos = np.asarray(xr_action["grip_pos"], dtype=float)
|
||||
grip_quat = np.asarray(xr_action["grip_quat"], dtype=float)
|
||||
squeeze = float(xr_action["squeeze"])
|
||||
trigger = float(xr_action["trigger"])
|
||||
enabled = squeeze > teleop_config.clutch_threshold
|
||||
|
||||
# On the engage edge, latch the clutch home at the arm's MEASURED EE pose (FK of
|
||||
# the live joints) and the controller origin so the per-frame delta starts at zero.
|
||||
# Latching the last commanded pose instead would snap the arm back to it at full
|
||||
# servo speed if the arm moved while disengaged (gravity sag, external contact).
|
||||
is_engage_frame = enabled and not prev_enabled
|
||||
if is_engage_frame:
|
||||
q_measured = np.array([float(robot_obs[f"{name}.pos"]) for name in motor_names], dtype=float)
|
||||
measured_base_T_ee = kinematics_solver.forward_kinematics(q_measured) # noqa: N806
|
||||
clutch.engage(grip_pos, grip_quat, measured_base_T_ee=measured_base_T_ee)
|
||||
# Re-anchor the pipeline state at the measured pose as well: EEBoundsAndSafety's
|
||||
# rate limiter and the IK warm start otherwise still reference the stale
|
||||
# pre-disengage command and would fight the fresh home for several frames.
|
||||
xr_to_robot_joints_processor.reset()
|
||||
prev_enabled = enabled
|
||||
|
||||
# SAFETY GATE: command the robot ONLY while the clutch is engaged; otherwise return
|
||||
# None so the loop holds the measured joints (releasing the clutch freezes the arm).
|
||||
if not enabled:
|
||||
return None
|
||||
|
||||
# Rebase the raw grip pose onto the EE, then run the pipeline. closedness = trigger.
|
||||
ee_pos, ee_quat = clutch.rebase(grip_pos, grip_quat)
|
||||
ee_action = {
|
||||
"ee_pose": np.concatenate([ee_pos, ee_quat]).astype(np.float32),
|
||||
"closedness": trigger,
|
||||
}
|
||||
return xr_to_robot_joints_processor((ee_action, robot_obs))
|
||||
|
||||
return Device(compute=compute, startup=startup, cleanup=teleop_device.disconnect)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# SO-101 leader arm device
|
||||
# ============================================================================
|
||||
|
||||
# Default duration [s] for the startup alignment slew (follower current -> leader first pose).
|
||||
ALIGN_DURATION_S = 3.0
|
||||
|
||||
# How long to wait for the leader plugin to start streaming before aligning / looping.
|
||||
LEADER_WARMUP_TIMEOUT_S = 20.0
|
||||
|
||||
# The plugin converts the leader's servo ticks to radians, so it reuses the serial SO-101
|
||||
# leader's calibration, stored by lerobot-calibrate under SO101Leader.name == "so_leader".
|
||||
SO_LEADER_CALIBRATION_NAME = "so_leader"
|
||||
|
||||
|
||||
def _leader_calibration_path(cfg: LoopConfig) -> Path | None:
|
||||
"""Infer the calibration JSON the launched plugin should read, or None.
|
||||
|
||||
Path convention: ``HF_LEROBOT_CALIBRATION / teleoperators / so_leader / {--teleop.id}.json``
|
||||
(or ``--teleop.calibration_dir`` if set). Returns None (plugin falls back to defaults) when
|
||||
it does not exist, warning if an id was given, or when no ``--teleop.id`` is set.
|
||||
"""
|
||||
if not cfg.teleop.id:
|
||||
return None
|
||||
calib_dir = cfg.teleop.calibration_dir or (
|
||||
HF_LEROBOT_CALIBRATION / TELEOPERATORS / SO_LEADER_CALIBRATION_NAME
|
||||
)
|
||||
calib_path = Path(calib_dir) / f"{cfg.teleop.id}.json"
|
||||
if calib_path.is_file():
|
||||
return calib_path
|
||||
print(
|
||||
f"WARNING: no leader calibration at {calib_path}; the plugin will use built-in defaults. "
|
||||
f"Calibrate with the serial leader (`lerobot-calibrate --teleop.type=so101_leader "
|
||||
f"--teleop.id={cfg.teleop.id}`) or the plugin's `calibrate` subcommand."
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _wait_for_leader(teleop: SO101LeaderArm, timeout_s: float) -> dict[str, float]:
|
||||
"""Poll the leader until it streams a live frame; return that frame's ``{joint}.pos``.
|
||||
|
||||
Raises ``SystemExit`` if no live frame arrives within ``timeout_s`` (plugin not pushing,
|
||||
wrong ``--teleop.collection_id``, or CloudXR not up).
|
||||
"""
|
||||
print(f"Waiting up to {timeout_s:.0f}s for the so101_leader plugin to stream…")
|
||||
deadline = time.time() + timeout_s
|
||||
while time.time() < deadline:
|
||||
action = teleop.get_action()
|
||||
if teleop.is_tracking:
|
||||
print("Leader is streaming.")
|
||||
return action
|
||||
time.sleep(1.0 / FPS)
|
||||
raise SystemExit(
|
||||
f"FAILED: leader did not stream within {timeout_s:.0f}s. Is the so101_leader plugin "
|
||||
"running and pushing (check --teleop.collection_id)? Is CloudXR up?"
|
||||
)
|
||||
|
||||
|
||||
def _maybe_launch_plugin(cfg: LoopConfig) -> subprocess.Popen | None:
|
||||
"""Spawn the so101_leader plugin if ``--launch_plugin <path>`` was given (after connect())."""
|
||||
if cfg.launch_plugin is None:
|
||||
return None
|
||||
if not Path(cfg.launch_plugin).exists():
|
||||
raise SystemExit(
|
||||
f"plugin binary not found: {cfg.launch_plugin} (build it in the IsaacTeleop repo first)"
|
||||
)
|
||||
leader_port = cfg.teleop.port # SO101LeaderArmConfig.port, forwarded to the plugin
|
||||
backend = f"leader on {leader_port}" if leader_port else "synthetic trajectory"
|
||||
print(f"launching plugin: {cfg.launch_plugin} ({backend})")
|
||||
# Positional args: [device_path] [collection_id] [calibration_file]. Empty device_path ->
|
||||
# synthetic backend. Calibration (only real hardware needs it) is appended when a port is set.
|
||||
argv = [cfg.launch_plugin, leader_port, cfg.teleop.collection_id]
|
||||
if leader_port:
|
||||
calib_path = _leader_calibration_path(cfg)
|
||||
if calib_path is not None:
|
||||
argv.append(str(calib_path))
|
||||
print(f" leader calibration: {calib_path}")
|
||||
# Spawned after connect() so it inherits the CloudXR runtime env (XR_RUNTIME_JSON, ...).
|
||||
proc = subprocess.Popen(argv)
|
||||
time.sleep(1.5) # let it create its OpenXR session and start pushing
|
||||
return proc
|
||||
|
||||
|
||||
def setup_leader(cfg: LoopConfig, robot, motor_names: list[str]) -> Device:
|
||||
"""Build the SO-101 leader arm device bundle (1:1 joint mirror)."""
|
||||
teleop_config = cfg.teleop # SO101LeaderArmConfig (selected via --teleop.type=so101_leader)
|
||||
teleop = SO101LeaderArm(teleop_config)
|
||||
|
||||
plugin_proc: subprocess.Popen | None = None
|
||||
|
||||
def startup() -> None:
|
||||
nonlocal plugin_proc
|
||||
# connect() auto-launches CloudXR (unless opted out); spawn the plugin AFTER so it
|
||||
# inherits the runtime env. The plugin is reaped in cleanup().
|
||||
teleop.connect()
|
||||
plugin_proc = _maybe_launch_plugin(cfg)
|
||||
|
||||
if not teleop.is_connected:
|
||||
raise ValueError("Teleop is not connected!")
|
||||
|
||||
# Block until the leader streams a live frame (clear error if it never does).
|
||||
_wait_for_leader(teleop, LEADER_WARMUP_TIMEOUT_S)
|
||||
|
||||
if cfg.align:
|
||||
print(f"Aligning follower to leader over {cfg.align_duration:.1f}s…")
|
||||
|
||||
# Re-read the live leader pose once per step so alpha=1 lands on its current pose
|
||||
# from a single coherent frame.
|
||||
def _leader_target() -> dict[str, float]:
|
||||
leader_now = teleop.get_action()
|
||||
return {name: float(leader_now[f"{name}.pos"]) for name in motor_names}
|
||||
|
||||
slew(robot, motor_names, _leader_target, cfg.align_duration)
|
||||
print("Alignment complete.")
|
||||
|
||||
print(
|
||||
"Starting joint-mirror loop. Back-drive the leader to teleoperate the follower… (Ctrl-C to stop)"
|
||||
)
|
||||
|
||||
def compute(robot_obs: RobotObservation | None) -> RobotAction | None:
|
||||
leader_action = teleop.get_action()
|
||||
# Hold the follower at its measured pose when the leader drops out (stale stream)
|
||||
# rather than commanding a possibly-old target.
|
||||
if not teleop.is_tracking:
|
||||
return None
|
||||
return leader_action
|
||||
|
||||
def cleanup() -> None:
|
||||
# A plugin-reaping failure must not skip the session disconnect (and vice versa
|
||||
# the disconnect runs after the plugin stops pushing on it).
|
||||
try:
|
||||
if plugin_proc is not None:
|
||||
plugin_proc.terminate()
|
||||
try:
|
||||
plugin_proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
plugin_proc.kill()
|
||||
finally:
|
||||
teleop.disconnect()
|
||||
|
||||
return Device(compute=compute, startup=startup, cleanup=cleanup)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Shared setup
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def build_device(cfg: LoopConfig) -> tuple:
|
||||
"""Connect the follower, build the selected Isaac device, and run its pre-loop startup.
|
||||
|
||||
Connects the follower FIRST (so the startup slew / clutch-home seed can read live joints),
|
||||
dispatches on ``--teleop.type``, then runs ``device.startup()`` before returning. On any
|
||||
failure after ``connect()`` the follower is disconnected so the connection never leaks.
|
||||
|
||||
Returns ``(robot, device, motor_names)``.
|
||||
"""
|
||||
# Default the CloudXR input profile to this example's default.env unless the user overrode
|
||||
# it via --teleop.cloudxr_env_file.
|
||||
if cfg.teleop.cloudxr_env_file is None:
|
||||
cfg.teleop.cloudxr_env_file = CLOUDXR_ENV_FILE
|
||||
|
||||
# SO-101/SO-100 only (both share the SO-101 URDF), reject other followers.
|
||||
supported_robots = {"so101_follower", "so100_follower"}
|
||||
if cfg.robot.type not in supported_robots:
|
||||
raise ValueError(
|
||||
f"This example only supports SO-101/SO-100 followers ({sorted(supported_robots)}), "
|
||||
f"but got --robot.type={cfg.robot.type}."
|
||||
)
|
||||
|
||||
# The degree-based pipeline relies on --robot.use_degrees (default True).
|
||||
robot = make_robot_from_config(cfg.robot)
|
||||
# Connect FIRST so the startup slew and clutch-home seed can read live joints.
|
||||
robot.connect()
|
||||
# Everything after connect() can fail; this runs outside the callers' try/finally, so
|
||||
# disconnect the follower on any failure to avoid leaking the connection.
|
||||
device: Device | None = None
|
||||
try:
|
||||
# Joint names in action order, read from {name}.pos action features (robot-agnostic).
|
||||
motor_names = [key.removesuffix(".pos") for key in robot.action_features if key.endswith(".pos")]
|
||||
|
||||
if isinstance(cfg.teleop, SO101LeaderArmConfig):
|
||||
device = setup_leader(cfg, robot, motor_names)
|
||||
else:
|
||||
device = setup_xr(cfg, robot, motor_names)
|
||||
|
||||
device.startup()
|
||||
except BaseException:
|
||||
# Reap a partially-started device, then always disconnect the follower.
|
||||
if device is not None:
|
||||
with suppress(Exception):
|
||||
device.cleanup()
|
||||
robot.disconnect()
|
||||
raise
|
||||
|
||||
return robot, device, motor_names
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Keyboard control
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def init_keyboard_listener():
|
||||
"""Recording shortcuts, terminal-first so they work over SSH.
|
||||
|
||||
Whenever stdin is a TTY we use the stdlib :class:`TerminalKeyListener` directly rather
|
||||
than upstream's pynput-first :func:`init_keyboard_listener`, whose global listener would
|
||||
capture the workstation console instead of this (often SSH) terminal. With no TTY we defer
|
||||
to upstream (pynput on a GUI, else headless no-op).
|
||||
"""
|
||||
if not (sys.stdin is not None and sys.stdin.isatty()):
|
||||
from lerobot.utils.keyboard_input import init_keyboard_listener as _upstream
|
||||
|
||||
return _upstream()
|
||||
|
||||
from lerobot.utils.keyboard_input import TerminalKeyListener, apply_recording_control
|
||||
|
||||
events = {"exit_early": False, "rerecord_episode": False, "stop_recording": False}
|
||||
|
||||
# n/r/q are the arrow/Esc equivalents that survive escape-sequence splitting over laggy
|
||||
# SSH/VNC links. Case-insensitive so Shift+letter still works.
|
||||
def on_key(name: str) -> None:
|
||||
key = name.lower()
|
||||
if key in ("right", "n"):
|
||||
apply_recording_control("right", events)
|
||||
elif key in ("left", "r"):
|
||||
apply_recording_control("left", events)
|
||||
elif key in ("esc", "q"):
|
||||
apply_recording_control("esc", events)
|
||||
|
||||
listener = TerminalKeyListener(on_key)
|
||||
listener.start()
|
||||
logging.info(
|
||||
"Keyboard control via terminal — keep this terminal focused: "
|
||||
"Right/n = end episode early, Left/r = re-record, Esc/q = stop."
|
||||
)
|
||||
return listener, events
|
||||
@@ -1,21 +0,0 @@
|
||||
# CloudXR device-profile overrides for the Isaac Teleop XR -> SO-101 example.
|
||||
#
|
||||
# Passed to isaacteleop's CloudXRLauncher as `env_config` (via
|
||||
# XRControllerConfig.cloudxr_env_file). Format: KEY=value, one per line; `#`
|
||||
# comments and blank lines ignored; $VARS / ~ expanded. See
|
||||
# isaacteleop/cloudxr/env_config.py::_load_env_file.
|
||||
#
|
||||
# Runtime-resolved keys (XR_RUNTIME_JSON, XRT_NO_STDIN, NV_CXR_RUNTIME_DIR,
|
||||
# NV_CXR_OUTPUT_DIR) are reserved and ignored if set here.
|
||||
|
||||
# Transport profile the runtime advertises (CloudXR default: auto-webrtc).
|
||||
# "Quest3" also covers the Pico 4. Other values: auto-native, AppleVisionPro.
|
||||
NV_DEVICE_PROFILE=Quest3
|
||||
|
||||
# Input device discovery channels (both default to true; pinned for clarity).
|
||||
NV_CXR_ENABLE_PUSH_DEVICES=true
|
||||
NV_CXR_ENABLE_TENSOR_DATA=true
|
||||
|
||||
# Runtime logs to ~/.cloudxr/logs — helps debug connection issues
|
||||
# (e.g. "Failed to get OpenXR system: -35").
|
||||
NV_CXR_FILE_LOGGING=true
|
||||
@@ -1,40 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 NVIDIA Corporation and The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""NVIDIA Isaac Teleop teleoperators for LeRobot.
|
||||
|
||||
Each input device is an :class:`IsaacTeleopTeleoperator` subclass: :class:`XRController`
|
||||
(XR/VR controller) and :class:`SO101LeaderArm` (back-drivable SO-101 leader arm) ship today.
|
||||
"""
|
||||
|
||||
from .base import IsaacTeleopTeleoperator
|
||||
from .clutch import Clutch
|
||||
from .config_isaac_teleop import IsaacTeleopConfig, SO101LeaderArmConfig, XRControllerConfig
|
||||
from .teleop_so101_leader_arm import SO101LeaderArm, leader_joints_to_robot_action
|
||||
from .teleop_xr_controller import XRController
|
||||
from .xr_controller_processor import MapXRControllerActionToRobotAction
|
||||
|
||||
__all__ = [
|
||||
"Clutch",
|
||||
"IsaacTeleopConfig",
|
||||
"IsaacTeleopTeleoperator",
|
||||
"MapXRControllerActionToRobotAction",
|
||||
"SO101LeaderArm",
|
||||
"SO101LeaderArmConfig",
|
||||
"XRController",
|
||||
"XRControllerConfig",
|
||||
"leader_joints_to_robot_action",
|
||||
]
|
||||
@@ -1,282 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 NVIDIA Corporation and The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Shared base for NVIDIA Isaac Teleop-backed LeRobot teleoperators.
|
||||
|
||||
Isaac Teleop is a multi-modal framework: a single ``TeleopSession`` can be driven by
|
||||
XR controllers, hand tracking, Manus gloves, etc. Each modality is a
|
||||
:class:`Teleoperator` subclass in its own ``teleop_<device>.py``.
|
||||
|
||||
:class:`IsaacTeleopTeleoperator` owns what those devices share — the session
|
||||
lifecycle, the per-step staleness/worker-health guard, and the no-op calibration
|
||||
tracking devices need. A concrete device implements :meth:`_build_pipeline` (its
|
||||
retargeting graph) and :meth:`get_action` (usually via :meth:`_step`).
|
||||
|
||||
``isaacteleop`` is an optional NVIDIA dependency (install instructions in the example's
|
||||
``README.md``); its imports are guarded behind an availability check at module top, so this
|
||||
module imports without it and constructing a device fails fast with install instructions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import abc
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from lerobot.teleoperators.teleoperator import Teleoperator
|
||||
from lerobot.utils.import_utils import is_package_available
|
||||
|
||||
from .config_isaac_teleop import IsaacTeleopConfig
|
||||
|
||||
_isaacteleop_available = is_package_available("isaacteleop")
|
||||
|
||||
if TYPE_CHECKING or _isaacteleop_available:
|
||||
from isaacteleop.cloudxr import CloudXRLauncher
|
||||
from isaacteleop.retargeting_engine.interface import (
|
||||
ExecutionEvents,
|
||||
ExecutionState,
|
||||
GraphExecutable,
|
||||
RetargeterIO,
|
||||
)
|
||||
from isaacteleop.teleop_session_manager import TeleopSession, TeleopSessionConfig
|
||||
else:
|
||||
CloudXRLauncher = None
|
||||
ExecutionEvents = None
|
||||
ExecutionState = None
|
||||
GraphExecutable = None
|
||||
RetargeterIO = None
|
||||
TeleopSession = None
|
||||
TeleopSessionConfig = None
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Gripper closedness [0, 1] -> SO-101 follower motor units [0, 100] (RANGE_0_100, 100 = OPEN).
|
||||
# Shared by the XR processor and leader device, which invert via ``pos = (1 - c) * SCALE``.
|
||||
_GRIPPER_MOTOR_SCALE = 100.0
|
||||
|
||||
|
||||
def _require_isaacteleop() -> None:
|
||||
"""Fail fast with install pointers when the optional ``isaacteleop`` package is missing."""
|
||||
if not _isaacteleop_available:
|
||||
raise ImportError(
|
||||
"The 'isaacteleop' package is required for Isaac Teleop devices but is not "
|
||||
"installed. See examples/isaac_teleop_to_so101/README.md for install instructions."
|
||||
)
|
||||
|
||||
|
||||
class IsaacTeleopTeleoperator(Teleoperator):
|
||||
"""Abstract base for teleoperators backed by an Isaac Teleop ``TeleopSession``.
|
||||
|
||||
Owns the session lifecycle and the per-step health guard; subclasses supply
|
||||
:meth:`_build_pipeline` and :meth:`get_action`.
|
||||
"""
|
||||
|
||||
config_class = IsaacTeleopConfig
|
||||
|
||||
def __init__(self, config: IsaacTeleopConfig):
|
||||
_require_isaacteleop()
|
||||
super().__init__(config)
|
||||
self.config: IsaacTeleopConfig = config
|
||||
self._session: TeleopSession | None = None
|
||||
self._cloudxr_launcher: CloudXRLauncher | None = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Pipeline construction (device override point)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@abc.abstractmethod
|
||||
def _build_pipeline(self) -> GraphExecutable:
|
||||
"""Build this device's retargeting pipeline (the ``GraphExecutable`` for
|
||||
``TeleopSessionConfig.pipeline``). Called once in :meth:`connect`; its output
|
||||
keys must match what :meth:`get_action` unpacks.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lifecycle (shared)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
return self._session is not None
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
return True # Tracking devices are self-calibrating.
|
||||
|
||||
def calibrate(self) -> None:
|
||||
pass
|
||||
|
||||
def configure(self) -> None:
|
||||
pass
|
||||
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
"""Auto-launch the CloudXR runtime (unless opted out) and open the session.
|
||||
|
||||
The CloudXR launch blocks ~30s and, on the first run, prompts on stdin for the
|
||||
EULA (accept once via ``python -m isaacteleop.cloudxr --accept-eula``). Opt out
|
||||
when CloudXR runs externally via ``config.auto_launch_cloudxr=False`` or
|
||||
``LEROBOT_CLOUDXR_SKIP_AUTOLAUNCH=1`` (env var wins).
|
||||
"""
|
||||
if self._session is not None:
|
||||
raise RuntimeError("Already connected. Call disconnect() first.")
|
||||
|
||||
self._ensure_cloudxr_runtime()
|
||||
|
||||
try:
|
||||
pipeline = self._build_pipeline()
|
||||
session_config = TeleopSessionConfig(app_name=self.config.app_name, pipeline=pipeline)
|
||||
self._session = TeleopSession(session_config)
|
||||
self._session.__enter__()
|
||||
except Exception:
|
||||
self._session = None
|
||||
try:
|
||||
self._stop_cloudxr_runtime()
|
||||
except Exception:
|
||||
logger.exception("Failed to stop CloudXR runtime during connect() rollback")
|
||||
raise
|
||||
logger.info("Isaac Teleop session started: %s", self.config.app_name)
|
||||
|
||||
def disconnect(self) -> None:
|
||||
try:
|
||||
if self._session is not None:
|
||||
# Null the handle BEFORE __exit__: even a failed session teardown must not
|
||||
# wedge the device as is_connected (blocking every later connect/disconnect).
|
||||
session = self._session
|
||||
self._session = None
|
||||
session.__exit__(None, None, None)
|
||||
logger.info("Isaac Teleop session ended")
|
||||
finally:
|
||||
# Reap the CloudXR runtime even if session teardown raised, and even if no
|
||||
# session was ever established (e.g. the launcher came up but session creation
|
||||
# failed before this point); a no-op when we never launched CloudXR (opt-out /
|
||||
# externally-owned runtime), so we never stop a runtime we don't own.
|
||||
self._stop_cloudxr_runtime()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# CloudXR runtime (shared)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _ensure_cloudxr_runtime(self) -> None:
|
||||
"""Auto-launch the CloudXR runtime once, unless opted out.
|
||||
|
||||
Idempotent (no-op once the launcher is up). ``LEROBOT_CLOUDXR_SKIP_AUTOLAUNCH``
|
||||
is checked first and wins over ``config.auto_launch_cloudxr``. Constructing
|
||||
:class:`CloudXRLauncher` mutates the process env (``XR_RUNTIME_JSON`` etc.) and
|
||||
blocks until the runtime is ready or raises :class:`RuntimeError`.
|
||||
"""
|
||||
if self._cloudxr_launcher is not None:
|
||||
return
|
||||
|
||||
if os.environ.get("LEROBOT_CLOUDXR_SKIP_AUTOLAUNCH", "").strip() == "1":
|
||||
logger.info(
|
||||
"LEROBOT_CLOUDXR_SKIP_AUTOLAUNCH=1 set; skipping CloudXR auto-launch "
|
||||
"(assuming CloudXR is already running externally)"
|
||||
)
|
||||
return
|
||||
|
||||
if not self.config.auto_launch_cloudxr:
|
||||
logger.info(
|
||||
"config.auto_launch_cloudxr is False; skipping CloudXR auto-launch "
|
||||
"(assuming CloudXR is already running externally)"
|
||||
)
|
||||
return
|
||||
|
||||
logger.info("Launching CloudXR runtime (first run may prompt for EULA and take ~30s)...")
|
||||
|
||||
self._cloudxr_launcher = CloudXRLauncher(
|
||||
install_dir=str(Path.home() / ".cloudxr"),
|
||||
env_config=self.config.cloudxr_env_file,
|
||||
accept_eula=False,
|
||||
)
|
||||
|
||||
def _stop_cloudxr_runtime(self) -> None:
|
||||
"""Stop the auto-launched CloudXR runtime, if any.
|
||||
|
||||
Clean stop nulls the handle. On :class:`RuntimeError` the handle is RETAINED so
|
||||
the launcher's ``atexit`` hook owns the retry — a later :meth:`connect` then
|
||||
treats the retained runtime as still up and will not relaunch.
|
||||
"""
|
||||
if self._cloudxr_launcher is None:
|
||||
return
|
||||
try:
|
||||
self._cloudxr_launcher.stop()
|
||||
except RuntimeError:
|
||||
logger.warning("CloudXR runtime could not be terminated; handle retained for atexit cleanup")
|
||||
else:
|
||||
self._cloudxr_launcher = None
|
||||
logger.info("CloudXR runtime stopped")
|
||||
|
||||
def send_feedback(self, feedback: dict[str, Any]) -> None:
|
||||
pass # Haptic feedback not yet implemented.
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Stepping (shared)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _running_events(self) -> ExecutionEvents:
|
||||
"""Constant ``RUNNING`` ``ExecutionEvents`` for a device with no clutch lifecycle.
|
||||
|
||||
Keeps the stream flowing; ``reset`` stays ``False``. A clutched device that needs
|
||||
a real lifecycle should build its own ``ExecutionEvents`` instead.
|
||||
"""
|
||||
return ExecutionEvents(execution_state=ExecutionState.RUNNING, reset=False)
|
||||
|
||||
def _step(
|
||||
self,
|
||||
*,
|
||||
execution_events: ExecutionEvents | None = None,
|
||||
external_inputs: Mapping[str, Any] | None = None,
|
||||
) -> RetargeterIO:
|
||||
"""Step the session once and return the raw pipeline outputs.
|
||||
|
||||
Applies the shared guard: re-raises a retargeting-worker exception and warns on a
|
||||
stale frame. Subclasses call this from :meth:`get_action`.
|
||||
|
||||
Args:
|
||||
execution_events: The ``ExecutionEvents`` driving the session this frame.
|
||||
Devices with a lifecycle (clutch) MUST pass this every frame — when
|
||||
``None``, ``TeleopSession.step`` auto-fires ``RUNNING`` (the clutch would
|
||||
latch immediately and never stop).
|
||||
external_inputs: Per-step inputs (e.g. a static ``base_T_anchor``) in the
|
||||
``{leaf_node_name: {output_port_name: TensorGroup}}`` shape ``step`` expects.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If not connected, or if the retargeting worker raised.
|
||||
"""
|
||||
if self._session is None:
|
||||
raise RuntimeError("Not connected. Call connect() first.")
|
||||
|
||||
result = self._session.step(
|
||||
execution_events=execution_events,
|
||||
external_inputs=external_inputs,
|
||||
)
|
||||
|
||||
info = self._session.last_step_info
|
||||
if info is not None:
|
||||
if info.worker_exception is not None:
|
||||
raise RuntimeError(
|
||||
"Isaac Teleop retargeting worker raised an exception"
|
||||
) from info.worker_exception
|
||||
if info.frame_deadline_miss:
|
||||
logger.warning(
|
||||
"Isaac Teleop frame deadline miss (returned_age_frames=%s)",
|
||||
info.returned_age_frames,
|
||||
)
|
||||
return result
|
||||
@@ -1,102 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 NVIDIA Corporation and The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Engage-relative clutch for the XR -> SO-101 teleop loop.
|
||||
|
||||
Turns the raw controller grip pose into an absolute base-frame EE target, so the XR
|
||||
device can stay a thin raw-pose reader. Pure numpy + the local ``Rotation`` helper (no
|
||||
``isaacteleop``), so it is unit-testable without the XR runtime.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
from lerobot.utils.rotation import Rotation
|
||||
|
||||
|
||||
class Clutch:
|
||||
"""Engage-relative clutch for both position AND orientation.
|
||||
|
||||
Latch an origin on engage, then track the base-frame delta from it, applied
|
||||
independently to position and orientation. State:
|
||||
|
||||
- ``_last_commanded_pos`` / ``_last_commanded_rot``: last commanded EE pose; held
|
||||
while disengaged so the arm freezes where it was left.
|
||||
- ``_home_pos`` / ``_home_rot``: latched on engage — the EE pose the delta applies to.
|
||||
The position comes from the arm's MEASURED pose when the caller provides it (so an
|
||||
arm that moved while disengaged is not snapped back to a stale command); the
|
||||
orientation always comes from the last commanded rotation (see NOTE below).
|
||||
- ``_origin_pos`` / ``_origin_rot``: latched on engage — the controller pose the delta
|
||||
is measured against.
|
||||
|
||||
Each engaged frame :meth:`rebase` returns::
|
||||
|
||||
pos = home_pos + (grip_pos - origin_pos) # 1:1 controller -> EE translation
|
||||
rot = (R_ctrl @ R_origin ^ -1) @ R_home # base-frame delta, left-composed
|
||||
|
||||
On the engage edge the output is exactly the home pose (no teleport). The orientation
|
||||
delta is left-composed (base frame), so hand rotation about base Z maps to EE rotation
|
||||
about base Z. A re-clutch latches a fresh home/origin.
|
||||
|
||||
NOTE: ``_home_rot`` is the last *commanded* orientation even when the measured pose is
|
||||
supplied: the 5-DOF SO-101 tracks orientation only softly, so its measured wrist
|
||||
orientation persistently differs from the command, and latching the measurement would
|
||||
inject that offset into the commanded signal on every re-clutch. Position has no such
|
||||
tracking gap, and there latching the measurement is what prevents the snap-back.
|
||||
"""
|
||||
|
||||
def __init__(self, home_base_T_ee: np.ndarray): # noqa: N803
|
||||
# Seed the held pose from the arm's measured startup EE pose so the first
|
||||
# engage latches home there (no jump on the first squeeze).
|
||||
home = np.asarray(home_base_T_ee, dtype=float)
|
||||
self._last_commanded_pos = home[:3, 3].copy()
|
||||
self._last_commanded_rot = Rotation.from_matrix(home[:3, :3])
|
||||
self._home_pos = self._last_commanded_pos.copy()
|
||||
self._home_rot = self._last_commanded_rot
|
||||
self._origin_pos = np.zeros(3, dtype=float)
|
||||
self._origin_rot = Rotation.from_quat(np.array([0.0, 0.0, 0.0, 1.0]))
|
||||
|
||||
def engage(
|
||||
self,
|
||||
grip_pos: np.ndarray,
|
||||
grip_quat: np.ndarray,
|
||||
measured_base_T_ee: np.ndarray | None = None, # noqa: N803
|
||||
) -> None:
|
||||
"""Latch the engage home (where the arm is now) and controller origin.
|
||||
|
||||
Pass ``measured_base_T_ee`` (FK of the measured joints) so the home POSITION is
|
||||
where the arm physically is — if the arm moved while disengaged (gravity sag,
|
||||
external contact), latching the stale last-commanded position would make the
|
||||
first engaged frame command a full-speed jump back to it. The home ORIENTATION
|
||||
always stays the last commanded one (see the class NOTE).
|
||||
"""
|
||||
if measured_base_T_ee is not None:
|
||||
self._home_pos = np.asarray(measured_base_T_ee, dtype=float)[:3, 3].copy()
|
||||
else:
|
||||
self._home_pos = self._last_commanded_pos.copy()
|
||||
self._home_rot = self._last_commanded_rot
|
||||
self._origin_pos = np.asarray(grip_pos, dtype=float).copy()
|
||||
self._origin_rot = Rotation.from_quat(np.asarray(grip_quat, dtype=float))
|
||||
|
||||
def rebase(self, grip_pos: np.ndarray, grip_quat: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Return the absolute base-frame EE target ``(pos [m], quat [xyzw])`` for this frame."""
|
||||
pos = self._home_pos + (np.asarray(grip_pos, dtype=float) - self._origin_pos)
|
||||
rot_ctrl = Rotation.from_quat(np.asarray(grip_quat, dtype=float))
|
||||
rot = (rot_ctrl * self._origin_rot.inv()) * self._home_rot
|
||||
self._last_commanded_pos = pos.copy()
|
||||
self._last_commanded_rot = rot
|
||||
return pos, rot.as_quat()
|
||||
@@ -1,135 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 NVIDIA Corporation and The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Configuration dataclasses for NVIDIA Isaac Teleop-backed teleoperators.
|
||||
|
||||
:class:`IsaacTeleopConfig` holds the shared fields; each device adds its own subclass
|
||||
(e.g. :class:`XRControllerConfig`, :class:`SO101LeaderArmConfig`).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import ClassVar
|
||||
|
||||
from lerobot.teleoperators.config import TeleoperatorConfig
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class IsaacTeleopConfig(TeleoperatorConfig):
|
||||
"""Shared config for all Isaac Teleop-backed teleoperators.
|
||||
|
||||
Uses its own draccus ``_choice_registry`` (decoupled from the global
|
||||
:class:`TeleoperatorConfig` one) so ``--teleop.type`` on a field typed
|
||||
``IsaacTeleopConfig`` resolves against ONLY the Isaac devices — letting them claim
|
||||
short names (``xr_controller``, ``so101_leader``) without colliding with the global
|
||||
registry. These devices are selected by the example scripts, not routed through
|
||||
``make_teleoperator_from_config``.
|
||||
"""
|
||||
|
||||
_choice_registry: ClassVar[dict] = {}
|
||||
|
||||
app_name: str = "LeTeleop"
|
||||
"""Application name for the OpenXR / Isaac Teleop session."""
|
||||
|
||||
auto_launch_cloudxr: bool = True
|
||||
"""Auto-launch the CloudXR runtime on :meth:`connect`. Set ``False`` (or export
|
||||
``LEROBOT_CLOUDXR_SKIP_AUTOLAUNCH=1``, which wins) when CloudXR runs externally.
|
||||
"""
|
||||
|
||||
cloudxr_env_file: str | None = None
|
||||
"""Optional CloudXR device-profile ``.env`` (an INPUT profile selecting the headset
|
||||
transport) passed to ``CloudXRLauncher``. ``None`` keeps the default auto-WebRTC profile.
|
||||
"""
|
||||
|
||||
|
||||
# Static rebase from the OpenXR controller anchor frame (X=Right, Y=Up, Z=Backward) into the
|
||||
# robot base frame (X=Forward, Y=Left, Z=Up). A proper rotation (det=+1): controller motion
|
||||
# forward -> robot +X, right -> robot -Y (i.e. rightward), up -> robot +Z.
|
||||
_DEFAULT_BASE_T_ANCHOR: list[list[float]] = [
|
||||
[0.0, 0.0, -1.0, 0.0],
|
||||
[-1.0, 0.0, 0.0, 0.0],
|
||||
[0.0, 1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 0.0, 1.0],
|
||||
]
|
||||
|
||||
|
||||
@IsaacTeleopConfig.register_subclass("xr_controller")
|
||||
@dataclass(kw_only=True)
|
||||
class XRControllerConfig(IsaacTeleopConfig):
|
||||
"""Config for Isaac Teleop XR (VR) controller teleoperation.
|
||||
|
||||
Exposes the raw base-frame grip pose, squeeze, and trigger via ``ControllersSource``.
|
||||
No retargeters: the clutch and gripper mapping live in the owning loop.
|
||||
"""
|
||||
|
||||
hand_side: str = "right"
|
||||
"""Which controller hand to use: ``"left"`` or ``"right"``. A plain ``str`` (validated in
|
||||
``__post_init__``) because draccus cannot decode ``Literal``-typed fields from the CLI."""
|
||||
|
||||
clutch_threshold: float = 0.5
|
||||
"""Squeeze value above which the owning loop's clutch engages (held-to-enable). The
|
||||
device reports only the raw squeeze; the threshold is applied by the loop."""
|
||||
|
||||
base_T_anchor: list[list[float]] = field( # noqa: N815 (frameA_T_frameB transform-matrix convention)
|
||||
# Fresh copy per instance: returning the module-level list itself would alias one
|
||||
# mutable matrix across every config.
|
||||
default_factory=lambda: [row.copy() for row in _DEFAULT_BASE_T_ANCHOR]
|
||||
)
|
||||
"""Static 4x4 [row-major] transform rebasing the OpenXR controller anchor frame into
|
||||
the robot base frame. Defaults to OpenXR (X=Right, Y=Up, Z=Backward) -> robot
|
||||
(X=Forward, Y=Left, Z=Up). Plain nested lists so the config stays serializable.
|
||||
"""
|
||||
|
||||
def __post_init__(self):
|
||||
if self.hand_side not in ("left", "right"):
|
||||
raise ValueError(f"hand_side must be 'left' or 'right', got {self.hand_side!r}")
|
||||
|
||||
|
||||
# Provisional gripper open/close endpoints [rad], normalizing the streamed gripper angle
|
||||
# into the follower's RANGE_0_100 jaw target. Derived from the so101_leader plugin README's
|
||||
# example calibration (home_ticks=2048, range 2000..3000; angle = (ticks-home)*2*pi/4096).
|
||||
_DEFAULT_GRIPPER_OPEN_RAD = -0.074
|
||||
_DEFAULT_GRIPPER_CLOSE_RAD = 1.460
|
||||
|
||||
|
||||
@IsaacTeleopConfig.register_subclass("so101_leader")
|
||||
@dataclass(kw_only=True)
|
||||
class SO101LeaderArmConfig(IsaacTeleopConfig):
|
||||
"""Config for an Isaac Teleop SO-101 *leader arm* (generic joint-space device).
|
||||
|
||||
Mirrors the leader's joint angles 1:1 onto a follower SO-101. The leader state is
|
||||
streamed in radians by the native ``so101_leader`` plugin and read via a
|
||||
``JointStateSource``; the device converts arm joints to degrees and the gripper to the
|
||||
follower's RANGE_0_100 jaw target (no IK/clutch/retargeter on the LeRobot side).
|
||||
"""
|
||||
|
||||
port: str = ""
|
||||
"""Serial port of the physical LEADER arm (e.g. ``/dev/ttyACM1``), forwarded to the
|
||||
plugin (which reads the servos) when the example launches it. Empty -> the plugin runs
|
||||
its synthetic trajectory."""
|
||||
|
||||
collection_id: str = "so101_leader"
|
||||
"""Tensor collection id the leader plugin pushes on; must match the running
|
||||
``so101_leader`` plugin (its second positional arg, default ``"so101_leader"``)."""
|
||||
|
||||
gripper_open_rad: float = _DEFAULT_GRIPPER_OPEN_RAD
|
||||
"""Leader gripper angle [rad] at fully OPEN -> follower jaw 100. Provisional default;
|
||||
set from the plugin's ``calibrate`` subcommand. See ``_DEFAULT_GRIPPER_OPEN_RAD``."""
|
||||
|
||||
gripper_close_rad: float = _DEFAULT_GRIPPER_CLOSE_RAD
|
||||
"""Leader gripper angle [rad] at fully CLOSED -> follower jaw 0. Provisional default;
|
||||
set from the plugin's ``calibrate`` subcommand. See ``_DEFAULT_GRIPPER_CLOSE_RAD``."""
|
||||
@@ -1,186 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 NVIDIA Corporation and The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""SO-101 leader-arm device for NVIDIA Isaac Teleop, exposed to LeRobot.
|
||||
|
||||
The leader is a back-drivable SO-101 whose six joint angles are streamed (in radians) by
|
||||
the native ``so101_leader`` plugin; this device reads them via a ``JointStateSource`` and
|
||||
converts them into follower-ready ``{joint}.pos``. Same kinematics as the follower, so it
|
||||
needs no retargeting — a 1:1 joint mirror, direct joint drive.
|
||||
|
||||
Units (converted in the device so the output is always follower-valid):
|
||||
|
||||
* arm joints: ``rad2deg`` — correct only if the leader's calibrated zero and the follower's
|
||||
homing map to the same physical zero (the standard same-hardware assumption).
|
||||
* gripper: normalized from ``[gripper_open_rad, gripper_close_rad]`` to RANGE_0_100.
|
||||
|
||||
``isaacteleop`` imports are guarded behind the availability flag so this module — and the
|
||||
pure :func:`leader_joints_to_robot_action` converter — import without it (construction
|
||||
fails fast via the base class).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
|
||||
from lerobot.types import RobotAction
|
||||
|
||||
from .base import _GRIPPER_MOTOR_SCALE, IsaacTeleopTeleoperator, _isaacteleop_available
|
||||
from .config_isaac_teleop import SO101LeaderArmConfig
|
||||
|
||||
if TYPE_CHECKING or _isaacteleop_available:
|
||||
from isaacteleop.retargeting_engine.deviceio_source_nodes import JointStateSource
|
||||
from isaacteleop.retargeting_engine.interface import OutputCombiner
|
||||
else:
|
||||
JointStateSource = None
|
||||
OutputCombiner = None
|
||||
|
||||
# Canonical SO-101 DOF names and order — matches the plugin stream and the follower's motor
|
||||
# order. Passed to the ``JointStateSource`` as its output layout; the source maps by name and
|
||||
# :func:`_joints_group_to_rad` reads back by name, so a layout mismatch can't mislabel a DOF.
|
||||
SO101_LEADER_JOINTS = [
|
||||
"shoulder_pan",
|
||||
"shoulder_lift",
|
||||
"elbow_flex",
|
||||
"wrist_flex",
|
||||
"wrist_roll",
|
||||
"gripper",
|
||||
]
|
||||
|
||||
|
||||
def leader_joints_to_robot_action(
|
||||
joints_rad: dict[str, float],
|
||||
*,
|
||||
gripper_joint: str,
|
||||
gripper_open_rad: float,
|
||||
gripper_close_rad: float,
|
||||
) -> RobotAction:
|
||||
"""Convert streamed leader joint angles [rad] to follower-ready ``{joint}.pos``.
|
||||
|
||||
Pure (no ``isaacteleop``, no I/O). Iteration follows ``joints_rad`` insertion order, so
|
||||
pass it in :data:`SO101_LEADER_JOINTS` order for a stable layout. Arm joints are
|
||||
converted ``rad2deg``; ``gripper_joint`` is normalized from
|
||||
``[gripper_open_rad, gripper_close_rad]`` to RANGE_0_100 (clipped).
|
||||
"""
|
||||
action: RobotAction = {}
|
||||
span = gripper_close_rad - gripper_open_rad
|
||||
for name, rad in joints_rad.items():
|
||||
if name == gripper_joint:
|
||||
# Closedness c=0 at open, c=1 at closed; invert to the follower's 100=open jaw.
|
||||
closedness = 0.0 if span == 0.0 else (rad - gripper_open_rad) / span
|
||||
closedness = min(1.0, max(0.0, closedness))
|
||||
action[f"{name}.pos"] = (1.0 - closedness) * _GRIPPER_MOTOR_SCALE
|
||||
else:
|
||||
action[f"{name}.pos"] = float(np.rad2deg(rad))
|
||||
return action
|
||||
|
||||
|
||||
def _joints_group_to_rad(joints) -> dict[str, float]:
|
||||
"""Read a ``JointStateSource`` output group into ``{joint_name: angle [rad]}``.
|
||||
|
||||
Pure (duck-typed on the group). The group is positional but each slot carries its joint
|
||||
name in ``group.group_type.types``; we key off those names (not a positional index) so a
|
||||
layout mismatch surfaces as a wrong/missing key here rather than a mislabeled DOF.
|
||||
"""
|
||||
names = [t.name for t in joints.group_type.types]
|
||||
return {name: float(joints[i]) for i, name in enumerate(names)}
|
||||
|
||||
|
||||
class SO101LeaderArm(IsaacTeleopTeleoperator):
|
||||
"""SO-101 leader-arm teleoperator (joint-space), direct joint mirror to the follower.
|
||||
|
||||
Reads the six joint angles off a single ``JointStateSource`` each frame; no retargeter,
|
||||
no clutch. When the leader is not streaming, :meth:`get_action` returns the held-last
|
||||
joints and :attr:`is_tracking` is ``False`` so the owning loop can hold the follower.
|
||||
"""
|
||||
|
||||
config_class = SO101LeaderArmConfig
|
||||
name = "isaac_teleop_so101_leader"
|
||||
|
||||
def __init__(self, config: SO101LeaderArmConfig):
|
||||
super().__init__(config)
|
||||
self.config: SO101LeaderArmConfig = config
|
||||
# Held-last joint angles [rad], seeded at zero (URDF/home pose) so the first frames
|
||||
# before the plugin starts pushing read as the home pose, not garbage.
|
||||
self._last_joints_rad: dict[str, float] = dict.fromkeys(SO101_LEADER_JOINTS, 0.0)
|
||||
# Whether the most recent get_action() read live leader data (vs held-last).
|
||||
self._is_tracking = False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Pipeline construction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_pipeline(self) -> OutputCombiner:
|
||||
"""Build the joint-mirror pipeline: a single ``JointStateSource`` leaf that converts
|
||||
the raw stream into a name-keyed joint group. No retargeter (shared kinematics)."""
|
||||
source = JointStateSource(
|
||||
name="so101_leader",
|
||||
collection_id=self.config.collection_id,
|
||||
joint_names=SO101_LEADER_JOINTS,
|
||||
)
|
||||
return OutputCombiner({"joints": source.output(JointStateSource.JOINTS)})
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Action features
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
# Matches the serial SOLeader's action features so this is a drop-in joint-space
|
||||
# leader: one float `{joint}.pos` per DOF, sendable straight to an SO-101 follower.
|
||||
return {f"{name}.pos": float for name in SO101_LEADER_JOINTS}
|
||||
|
||||
@property
|
||||
def feedback_features(self) -> dict[str, type]:
|
||||
return {}
|
||||
|
||||
@property
|
||||
def is_tracking(self) -> bool:
|
||||
"""Whether the last :meth:`get_action` read live leader data (vs held-last)."""
|
||||
return self._is_tracking
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Action extraction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_action(self) -> RobotAction:
|
||||
"""Step the session and return the leader joints as follower-ready ``{joint}.pos``.
|
||||
|
||||
When the leader is streaming, the live angles are cached and converted; otherwise the
|
||||
held-last angles are reused and :attr:`is_tracking` is set ``False``.
|
||||
"""
|
||||
result = self._step(execution_events=self._running_events())
|
||||
|
||||
joints = result["joints"]
|
||||
# The JointStateSource output is Optional: absent (is_none) when the device is
|
||||
# inactive. Treat that as "not tracking" and reuse the held-last angles.
|
||||
self._is_tracking = not getattr(joints, "is_none", False)
|
||||
if self._is_tracking:
|
||||
try:
|
||||
self._last_joints_rad = _joints_group_to_rad(joints)
|
||||
except (AttributeError, IndexError, KeyError, TypeError, ValueError):
|
||||
# A partially-populated / malformed group on an odd frame: keep held-last, but
|
||||
# report it as not-tracking so the loop holds the follower rather than trusting it.
|
||||
self._is_tracking = False
|
||||
|
||||
return leader_joints_to_robot_action(
|
||||
self._last_joints_rad,
|
||||
gripper_joint="gripper",
|
||||
gripper_open_rad=self.config.gripper_open_rad,
|
||||
gripper_close_rad=self.config.gripper_close_rad,
|
||||
)
|
||||
@@ -1,204 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 NVIDIA Corporation and The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""XR (VR) controller device for NVIDIA Isaac Teleop, exposed to LeRobot.
|
||||
|
||||
A deliberately thin reader: exposes the raw controller grip pose off
|
||||
``ControllersSource`` (statically rebased into the robot base frame by
|
||||
``ControllerTransform``), plus squeeze and trigger. No retargeters and no clutch —
|
||||
the clutch rebasing and gripper mapping live downstream in the owning loop, so this
|
||||
device is stateless across frames.
|
||||
|
||||
``isaacteleop`` imports are guarded behind the availability flag so this module imports
|
||||
without it (construction fails fast via the base class).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from lerobot.types import RobotAction
|
||||
|
||||
from .base import IsaacTeleopTeleoperator, _isaacteleop_available
|
||||
from .config_isaac_teleop import XRControllerConfig
|
||||
|
||||
if TYPE_CHECKING or _isaacteleop_available:
|
||||
from isaacteleop.retargeting_engine.deviceio_source_nodes import ControllersSource
|
||||
from isaacteleop.retargeting_engine.interface import OutputCombiner, TensorGroup, ValueInput
|
||||
from isaacteleop.retargeting_engine.tensor_types import TransformMatrix
|
||||
from isaacteleop.retargeting_engine.tensor_types.indices import ControllerInputIndex
|
||||
else:
|
||||
ControllersSource = None
|
||||
OutputCombiner = None
|
||||
TensorGroup = None
|
||||
ValueInput = None
|
||||
TransformMatrix = None
|
||||
ControllerInputIndex = None
|
||||
|
||||
# Source-node name for the static base_T_anchor rebase input fed via
|
||||
# ``TeleopSession.step(external_inputs=...)`` each frame.
|
||||
_BASE_T_ANCHOR_INPUT = "base_T_anchor"
|
||||
|
||||
|
||||
class XRController(IsaacTeleopTeleoperator):
|
||||
"""Raw XR controller grip-pose teleoperator (base-frame), no retargeters.
|
||||
|
||||
Reads the raw grip pose + squeeze + trigger off a ``ControllersSource`` rebased into
|
||||
the robot base frame. :meth:`get_action` returns the absolute base-frame grip pose
|
||||
untouched; the owning loop owns the clutch and gripper mapping.
|
||||
"""
|
||||
|
||||
config_class = XRControllerConfig
|
||||
name = "isaac_teleop_controller"
|
||||
|
||||
def __init__(self, config: XRControllerConfig):
|
||||
super().__init__(config)
|
||||
self.config: XRControllerConfig = config
|
||||
|
||||
# Constant base_T_anchor input, built once in connect() (a TensorGroup is heavy and
|
||||
# isaacteleop-backed) and reused every step.
|
||||
self._external_inputs: dict[str, Any] | None = None
|
||||
# Whether the last get_action() read a tracked controller; the owning loop polls this
|
||||
# to wait for the operator to connect before driving the arm.
|
||||
self._is_tracking = False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Pipeline construction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_pipeline(self) -> OutputCombiner:
|
||||
"""Build the raw-grip-pose pipeline: a ``ControllersSource`` rebased into the base
|
||||
frame by ``ControllerTransform``, exposed verbatim as ``"controller"``. No retargeters.
|
||||
"""
|
||||
side = self.config.hand_side
|
||||
controller_key = f"controller_{side}"
|
||||
|
||||
controllers = ControllersSource(name="controllers")
|
||||
# Static base_T_anchor rebase fed via external_inputs each step.
|
||||
xform = ValueInput(_BASE_T_ANCHOR_INPUT, TransformMatrix())
|
||||
transformed = controllers.transformed(xform.output("value"))
|
||||
ctrl = transformed.output(controller_key)
|
||||
|
||||
return OutputCombiner({"controller": ctrl})
|
||||
|
||||
def _build_external_inputs(self) -> dict[str, Any]:
|
||||
"""Materialize the constant ``base_T_anchor`` external input (once, in connect)."""
|
||||
tg = TensorGroup(TransformMatrix())
|
||||
tg[0] = np.asarray(self.config.base_T_anchor, dtype=np.float32)
|
||||
return {_BASE_T_ANCHOR_INPUT: {"value": tg}}
|
||||
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
super().connect(calibrate=calibrate)
|
||||
try:
|
||||
self._external_inputs = self._build_external_inputs()
|
||||
except Exception:
|
||||
# Roll the session/runtime back so a failed connect() leaves no half-state
|
||||
# (a live session behind a raised connect would leak the CloudXR runtime).
|
||||
self.disconnect()
|
||||
raise
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Action features
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def action_features(self) -> dict:
|
||||
return {
|
||||
"grip_pos": {
|
||||
"dtype": "float32",
|
||||
"shape": (3,),
|
||||
"names": {"x": 0, "y": 1, "z": 2},
|
||||
},
|
||||
"grip_quat": {
|
||||
"dtype": "float32",
|
||||
"shape": (4,),
|
||||
"names": {"qx": 0, "qy": 1, "qz": 2, "qw": 3},
|
||||
},
|
||||
# ``get_action`` returns scalars for these two, so the advertised
|
||||
# shape is () (0-d) to stay consistent with the returned values.
|
||||
"squeeze": {
|
||||
"dtype": "float32",
|
||||
"shape": (),
|
||||
"names": None,
|
||||
},
|
||||
"trigger": {
|
||||
"dtype": "float32",
|
||||
"shape": (),
|
||||
"names": None,
|
||||
},
|
||||
}
|
||||
|
||||
@property
|
||||
def feedback_features(self) -> dict:
|
||||
return {}
|
||||
|
||||
@property
|
||||
def is_tracking(self) -> bool:
|
||||
"""Whether the last :meth:`get_action` read a tracked controller. ``False`` until the
|
||||
headset is connected over CloudXR and its controllers are live; the owning loop polls
|
||||
it to wait for the operator before commanding the arm."""
|
||||
return self._is_tracking
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Action extraction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_action(self) -> RobotAction:
|
||||
"""Step the session and return the raw base-frame grip pose.
|
||||
|
||||
Reads the grip pose + squeeze + trigger off the transformed controller stream (with
|
||||
the constant ``base_T_anchor`` rebase). When the controller is not tracked, returns
|
||||
identity pose and squeeze/trigger = 0.0 so the owning loop freezes the arm.
|
||||
|
||||
Returns:
|
||||
``{"grip_pos": (3,) [m], "grip_quat": (4,) [qx,qy,qz,qw], "squeeze": float,
|
||||
"trigger": float}`` — pose in the robot base frame; squeeze/trigger in ``[0, 1]``.
|
||||
"""
|
||||
result = self._step(execution_events=self._running_events(), external_inputs=self._external_inputs)
|
||||
|
||||
# Optional controller group is None until the headset is connected and its controllers
|
||||
# are live; expose that as is_tracking so the loop can wait before driving the arm.
|
||||
controller = result["controller"]
|
||||
grip_pos = np.zeros(3, dtype=np.float32)
|
||||
grip_quat = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32)
|
||||
squeeze = 0.0
|
||||
trigger = 0.0
|
||||
self._is_tracking = not getattr(controller, "is_none", False)
|
||||
if self._is_tracking:
|
||||
# Read ALL four fields into locals before committing any of them: a failure on a
|
||||
# partially-populated frame must not mix live values with the safe defaults (a
|
||||
# live squeeze paired with a defaulted trigger=0.0 would keep the clutch engaged
|
||||
# while commanding the gripper fully open, dropping whatever is grasped). On
|
||||
# failure the defaults stand untouched and the frame reports not-tracked.
|
||||
try:
|
||||
pos = np.asarray(controller[ControllerInputIndex.GRIP_POSITION], dtype=np.float32)
|
||||
quat = np.asarray(controller[ControllerInputIndex.GRIP_ORIENTATION], dtype=np.float32)
|
||||
squeeze_val = float(controller[ControllerInputIndex.SQUEEZE_VALUE])
|
||||
trigger_val = float(controller[ControllerInputIndex.TRIGGER_VALUE])
|
||||
except (IndexError, KeyError, TypeError, ValueError):
|
||||
self._is_tracking = False
|
||||
else:
|
||||
grip_pos, grip_quat = pos, quat
|
||||
squeeze, trigger = squeeze_val, trigger_val
|
||||
|
||||
return {
|
||||
"grip_pos": grip_pos,
|
||||
"grip_quat": grip_quat,
|
||||
"squeeze": squeeze,
|
||||
"trigger": trigger,
|
||||
}
|
||||
@@ -1,87 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 NVIDIA Corporation and The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Processor step that maps XR controller actions to robot EE targets.
|
||||
|
||||
Analogous to ``MapPhoneActionToRobotAction``, this bridges the clutch-rebased EE pose to
|
||||
the IK pipeline's input contract (``EEBoundsAndSafety`` -> ``InverseKinematicsEEToJoints``).
|
||||
Pure (no ``isaacteleop``), so it is unit-testable without the XR runtime.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import ProcessorStepRegistry, RobotActionProcessorStep
|
||||
from lerobot.types import RobotAction
|
||||
from lerobot.utils.rotation import Rotation
|
||||
|
||||
from .base import _GRIPPER_MOTOR_SCALE
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register("map_xr_controller_action_to_robot_action")
|
||||
@dataclass
|
||||
class MapXRControllerActionToRobotAction(RobotActionProcessorStep):
|
||||
"""Maps an absolute base-frame EE pose + gripper closedness to the IK input contract.
|
||||
|
||||
Pure, stateless rename (the owning loop's clutch already produced the absolute base-frame
|
||||
target). Each frame it writes:
|
||||
|
||||
- ``ee.x/y/z`` = ``ee_pose[:3]`` (position [m]);
|
||||
- ``ee.wx/wy/wz`` = rotvec of ``ee_pose[3:7]`` (orientation; the IK tracks it softly at a
|
||||
small ``orientation_weight`` on the 5-DOF SO-101);
|
||||
- ``ee.gripper_pos`` = ``(1 - closedness) * _GRIPPER_MOTOR_SCALE`` (jaw target [0, 100],
|
||||
RANGE_0_100 where 100 = open, so closedness is inverted).
|
||||
|
||||
Input keys: ``ee_pose`` ``(7,)`` ``[x,y,z,qx,qy,qz,qw]``, ``closedness`` float in [0, 1].
|
||||
"""
|
||||
|
||||
def action(self, action: RobotAction) -> RobotAction:
|
||||
ee_pose = action.pop("ee_pose")
|
||||
closedness = float(action.pop("closedness"))
|
||||
|
||||
action["ee.x"] = float(ee_pose[0])
|
||||
action["ee.y"] = float(ee_pose[1])
|
||||
action["ee.z"] = float(ee_pose[2])
|
||||
# Orientation target as a rotvec (quat [qx,qy,qz,qw] -> axis-angle); the IK
|
||||
# consumes ee.w* as a rotvec and tracks it with orientation_weight.
|
||||
rotvec = Rotation.from_quat(ee_pose[3:7]).as_rotvec()
|
||||
action["ee.wx"] = float(rotvec[0])
|
||||
action["ee.wy"] = float(rotvec[1])
|
||||
action["ee.wz"] = float(rotvec[2])
|
||||
# Inverted: closedness c=1 (closed) -> 0, c=0 (open) -> 100 (SO-101 calibration).
|
||||
action["ee.gripper_pos"] = (1.0 - closedness) * _GRIPPER_MOTOR_SCALE
|
||||
return action
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
for feat in ["ee_pose", "closedness"]:
|
||||
features[PipelineFeatureType.ACTION].pop(feat, None)
|
||||
|
||||
for feat in [
|
||||
"ee.x",
|
||||
"ee.y",
|
||||
"ee.z",
|
||||
"ee.wx",
|
||||
"ee.wy",
|
||||
"ee.wz",
|
||||
"ee.gripper_pos",
|
||||
]:
|
||||
features[PipelineFeatureType.ACTION][feat] = PolicyFeature(type=FeatureType.ACTION, shape=(1,))
|
||||
|
||||
return features
|
||||
@@ -1,73 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 NVIDIA Corporation and The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Save the current SO-101 joint positions as the reset-origin pose (override).
|
||||
|
||||
Move the arm to the desired reset pose by hand (torque off), then run this script to write
|
||||
those joints to a per-arm file in the LeRobot cache. ``teleoperate.py`` / ``record.py`` load
|
||||
it on startup (matched by ``--robot.id``) as the reset target instead of the defaults.
|
||||
|
||||
Usage::
|
||||
|
||||
# 1. Move arm to desired reset pose by hand
|
||||
python -m examples.isaac_teleop_to_so101.override_reset_pose [--port /dev/ttyACM0] [--id so101_follower_arm]
|
||||
|
||||
# 2. Launch teleop with the SAME --robot.id — it will now reset to this pose on startup
|
||||
python -m examples.isaac_teleop_to_so101.teleoperate --robot.type=so101_follower --robot.port=/dev/ttyACM0 --robot.id=so101_follower_arm --teleop.type=xr_controller
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
|
||||
|
||||
from .common import RESET_POSE_FILE
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
||||
)
|
||||
parser.add_argument("--port", type=str, default="/dev/ttyACM0")
|
||||
parser.add_argument("--id", type=str, default="so101_follower_arm")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
robot = SO100Follower(SO100FollowerConfig(port=args.port, id=args.id, use_degrees=True))
|
||||
robot.connect()
|
||||
# Always disconnect the follower so a failure never leaks the serial connection.
|
||||
try:
|
||||
obs = robot.get_observation()
|
||||
motor_names = list(robot.bus.motors.keys())
|
||||
pose = {name: float(obs[f"{name}.pos"]) for name in motor_names}
|
||||
finally:
|
||||
robot.disconnect()
|
||||
|
||||
print("Current joint positions:")
|
||||
for name, val in pose.items():
|
||||
print(f" {name:20s}: {val:.2f}")
|
||||
|
||||
reset_pose_file = Path(RESET_POSE_FILE.format(robot_name=robot.name, robot_id=robot.id))
|
||||
reset_pose_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
reset_pose_file.write_text(json.dumps(pose, indent=2))
|
||||
print(f"\nSaved to {reset_pose_file}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,321 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 NVIDIA Corporation and The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Record a LeRobot dataset via NVIDIA Isaac Teleop -> SO-101.
|
||||
|
||||
Runs ``teleoperate.py``'s control loop while also saving each frame to a LeRobot dataset.
|
||||
``--teleop.type`` selects the device (``xr_controller`` | ``so101_leader``) as in
|
||||
``teleoperate.py``.
|
||||
|
||||
Usage::
|
||||
|
||||
# XR (VR) controller: clutch + soft-orientation IK
|
||||
python -m examples.isaac_teleop_to_so101.record \\
|
||||
--robot.type=so101_follower \\
|
||||
--robot.port=/dev/ttyACM0 \\
|
||||
--robot.id=so101_follower_arm \\
|
||||
--teleop.type=xr_controller \\
|
||||
--robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \\
|
||||
--dataset.repo_id=<hf_user>/<dataset_name> \\
|
||||
--dataset.single_task="Pick up vial from rack on the left side" \\
|
||||
--dataset.num_episodes=3 \\
|
||||
--dataset.episode_time_s=20 \\
|
||||
--dataset.reset_time_s=5
|
||||
|
||||
# SO-101 leader arm: 1:1 joint mirror (real leader on /dev/ttyACM1)
|
||||
python -m examples.isaac_teleop_to_so101.record \\
|
||||
--robot.type=so101_follower --robot.port=/dev/ttyACM0 --robot.id=so101_follower_arm \\
|
||||
--teleop.type=so101_leader --teleop.port=/dev/ttyACM1 --teleop.id=so101_leader_arm \\
|
||||
--launch_plugin=/path/to/IsaacTeleop/install/plugins/so101_leader/so101_leader_plugin \\
|
||||
--dataset.repo_id=<hf_user>/<dataset_name> --dataset.single_task="Pick up the cube" \\
|
||||
--dataset.num_episodes=3 --dataset.episode_time_s=20 --dataset.reset_time_s=5
|
||||
|
||||
The loop/launch knobs mirror ``teleoperate.py`` (tagged ``[xr]`` / ``[leader]`` below).
|
||||
|
||||
Keyboard shortcuts: Right/n = end episode early and save, Left/r = discard + re-record,
|
||||
Esc/q = stop after the current episode. All frames are recorded (including hold frames).
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import asdict, dataclass
|
||||
from pprint import pformat
|
||||
|
||||
from lerobot.cameras import CameraConfig # noqa: F401
|
||||
from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401
|
||||
from lerobot.common.control_utils import sanity_check_dataset_robot_compatibility
|
||||
from lerobot.configs import parser
|
||||
from lerobot.configs.dataset import DatasetRecordConfig
|
||||
from lerobot.datasets import (
|
||||
LeRobotDataset,
|
||||
VideoEncodingManager,
|
||||
aggregate_pipeline_dataset_features,
|
||||
create_initial_features,
|
||||
safe_stop_image_writer,
|
||||
)
|
||||
from lerobot.processor import make_default_processors
|
||||
from lerobot.robots import RobotConfig
|
||||
from lerobot.robots.so_follower import SOFollowerConfig # noqa: F401 (registers so101_follower)
|
||||
from lerobot.utils.constants import ACTION, OBS_STR
|
||||
from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts
|
||||
from lerobot.utils.robot_utils import precise_sleep
|
||||
from lerobot.utils.utils import init_logging
|
||||
|
||||
from .common import (
|
||||
ALIGN_DURATION_S,
|
||||
RESET_DURATION_S,
|
||||
Device,
|
||||
HoldLatch,
|
||||
build_device,
|
||||
init_keyboard_listener,
|
||||
)
|
||||
from .isaac_teleop import IsaacTeleopConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class RecordConfig:
|
||||
"""CLI config for Isaac Teleop -> SO-101 dataset recording.
|
||||
|
||||
``--robot.*`` / ``--teleop.*`` / ``--dataset.*`` configure the follower, device, and
|
||||
recording; the loop/launch knobs below carry the same ``[xr]`` / ``[leader]`` tags as
|
||||
``teleoperate.py``. Use ``--flag=false`` for booleans (draccus style).
|
||||
"""
|
||||
|
||||
robot: RobotConfig
|
||||
# --teleop.type=xr_controller|so101_leader, resolved against IsaacTeleopConfig's registry.
|
||||
teleop: IsaacTeleopConfig
|
||||
dataset: DatasetRecordConfig
|
||||
|
||||
# [leader] Path to the so101_leader plugin binary to spawn after CloudXR is up (it then
|
||||
# inherits the runtime env). None (default) -> assume the plugin already runs externally.
|
||||
launch_plugin: str | None = None
|
||||
|
||||
# [xr] Slew all joints to the reset pose before the first episode (--reset_to_origin=false to
|
||||
# keep the arm where it is). After the slew the clutch seeds its home from the measured pose.
|
||||
reset_to_origin: bool = True
|
||||
# [xr] Duration [s] of the reset-to-origin slew (passed through to setup_xr).
|
||||
reset_duration: float = RESET_DURATION_S
|
||||
|
||||
# [leader] Slew the follower to the leader's first pose before mirroring (--align=false to
|
||||
# begin the 1:1 mirror immediately; the follower may snap).
|
||||
align: bool = True
|
||||
# [leader] Duration [s] of the startup alignment slew.
|
||||
align_duration: float = ALIGN_DURATION_S
|
||||
|
||||
# Resume recording on an existing (previously interrupted) dataset.
|
||||
resume: bool = False
|
||||
|
||||
|
||||
@safe_stop_image_writer
|
||||
def _record_loop(
|
||||
robot,
|
||||
device: Device,
|
||||
motor_names: list[str],
|
||||
events: dict,
|
||||
fps: int,
|
||||
dataset: LeRobotDataset | None = None,
|
||||
control_time_s: float = 0.0,
|
||||
single_task: str | None = None,
|
||||
) -> None:
|
||||
"""Run one episode (or reset phase) of the control loop.
|
||||
|
||||
When ``dataset`` is None the loop still controls the robot (so the operator
|
||||
can reposition the arm during the reset window) but does not record frames.
|
||||
"""
|
||||
control_interval = 1.0 / fps
|
||||
timestamp = 0.0
|
||||
start_t = time.perf_counter()
|
||||
record_frames = dataset is not None
|
||||
hold = HoldLatch(motor_names)
|
||||
|
||||
while timestamp < control_time_s:
|
||||
loop_start = time.perf_counter()
|
||||
|
||||
if events["exit_early"]:
|
||||
events["exit_early"] = False
|
||||
break
|
||||
|
||||
obs = robot.get_observation()
|
||||
|
||||
if record_frames:
|
||||
observation_frame = build_dataset_frame(dataset.features, obs, prefix=OBS_STR)
|
||||
|
||||
# Device idle (XR clutch disengaged, or leader stream stale) -> hold the pose
|
||||
# latched on the active->idle edge.
|
||||
action = hold.resolve(device.compute(obs), obs)
|
||||
|
||||
robot.send_action(action)
|
||||
|
||||
if record_frames:
|
||||
action_frame = build_dataset_frame(dataset.features, action, prefix=ACTION)
|
||||
dataset.add_frame({**observation_frame, **action_frame, "task": single_task})
|
||||
|
||||
dt_s = time.perf_counter() - loop_start
|
||||
precise_sleep(max(control_interval - dt_s, 0.0))
|
||||
timestamp = time.perf_counter() - start_t
|
||||
|
||||
|
||||
@parser.wrap()
|
||||
def record(cfg: RecordConfig) -> LeRobotDataset:
|
||||
init_logging()
|
||||
logging.info(pformat(asdict(cfg)))
|
||||
|
||||
# Connect the follower, build the selected Isaac device, and run its pre-loop startup
|
||||
# (reset slew / leader align) — shared with teleoperate.py.
|
||||
robot, device, motor_names = build_device(cfg)
|
||||
|
||||
# Build dataset feature spec. The IK pipeline lives inside device.compute(), so the
|
||||
# action features are exactly robot.action_features (joint positions in degrees).
|
||||
teleop_proc, _, obs_proc = make_default_processors()
|
||||
dataset_features = combine_feature_dicts(
|
||||
aggregate_pipeline_dataset_features(
|
||||
pipeline=teleop_proc,
|
||||
initial_features=create_initial_features(action=robot.action_features),
|
||||
use_videos=cfg.dataset.video,
|
||||
),
|
||||
aggregate_pipeline_dataset_features(
|
||||
pipeline=obs_proc,
|
||||
initial_features=create_initial_features(observation=robot.observation_features),
|
||||
use_videos=cfg.dataset.video,
|
||||
),
|
||||
)
|
||||
|
||||
num_cameras = len(robot.cameras) if hasattr(robot, "cameras") else 0
|
||||
image_writer_threads = cfg.dataset.num_image_writer_threads_per_camera * num_cameras
|
||||
|
||||
dataset: LeRobotDataset | None = None
|
||||
listener = None
|
||||
try:
|
||||
if cfg.resume:
|
||||
dataset = LeRobotDataset.resume(
|
||||
cfg.dataset.repo_id,
|
||||
root=cfg.dataset.root,
|
||||
batch_encoding_size=cfg.dataset.video_encoding_batch_size,
|
||||
rgb_encoder=cfg.dataset.rgb_encoder,
|
||||
depth_encoder=cfg.dataset.depth_encoder,
|
||||
encoder_threads=cfg.dataset.encoder_threads,
|
||||
streaming_encoding=cfg.dataset.streaming_encoding,
|
||||
encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,
|
||||
image_writer_processes=cfg.dataset.num_image_writer_processes if num_cameras > 0 else 0,
|
||||
image_writer_threads=image_writer_threads if num_cameras > 0 else 0,
|
||||
)
|
||||
sanity_check_dataset_robot_compatibility(dataset, robot, cfg.dataset.fps, dataset_features)
|
||||
else:
|
||||
cfg.dataset.stamp_repo_id()
|
||||
dataset = LeRobotDataset.create(
|
||||
cfg.dataset.repo_id,
|
||||
cfg.dataset.fps,
|
||||
root=cfg.dataset.root,
|
||||
robot_type=robot.name,
|
||||
features=dataset_features,
|
||||
use_videos=cfg.dataset.video,
|
||||
image_writer_processes=cfg.dataset.num_image_writer_processes,
|
||||
image_writer_threads=image_writer_threads,
|
||||
batch_encoding_size=cfg.dataset.video_encoding_batch_size,
|
||||
rgb_encoder=cfg.dataset.rgb_encoder,
|
||||
depth_encoder=cfg.dataset.depth_encoder,
|
||||
encoder_threads=cfg.dataset.encoder_threads,
|
||||
streaming_encoding=cfg.dataset.streaming_encoding,
|
||||
encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,
|
||||
)
|
||||
|
||||
listener, events = init_keyboard_listener()
|
||||
|
||||
loop_kwargs = {
|
||||
"robot": robot,
|
||||
"device": device,
|
||||
"motor_names": motor_names,
|
||||
"events": events,
|
||||
"fps": cfg.dataset.fps,
|
||||
"single_task": cfg.dataset.single_task,
|
||||
}
|
||||
|
||||
with VideoEncodingManager(dataset):
|
||||
recorded_episodes = 0
|
||||
while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]:
|
||||
logging.info(f"Recording episode {dataset.num_episodes}")
|
||||
_record_loop(
|
||||
**loop_kwargs,
|
||||
dataset=dataset,
|
||||
control_time_s=cfg.dataset.episode_time_s,
|
||||
)
|
||||
|
||||
# Reset window: give the operator time to reposition the scene.
|
||||
# Skipped for the last episode (or if stop_recording was set).
|
||||
if not events["stop_recording"] and (
|
||||
recorded_episodes < cfg.dataset.num_episodes - 1 or events["rerecord_episode"]
|
||||
):
|
||||
logging.info("Reset the environment")
|
||||
_record_loop(
|
||||
**loop_kwargs,
|
||||
dataset=None,
|
||||
control_time_s=cfg.dataset.reset_time_s,
|
||||
)
|
||||
|
||||
if events["rerecord_episode"]:
|
||||
logging.info("Re-record episode")
|
||||
events["rerecord_episode"] = False
|
||||
events["exit_early"] = False
|
||||
dataset.clear_episode_buffer()
|
||||
continue
|
||||
|
||||
dataset.save_episode()
|
||||
recorded_episodes += 1
|
||||
|
||||
finally:
|
||||
logging.info("Stop recording")
|
||||
|
||||
# Hardware teardown FIRST, each step guarded: the arm must be freed promptly (not
|
||||
# after a potentially long finalize/encode), a cleanup failure must not skip the
|
||||
# follower disconnect (which is what disables torque), and neither must prevent
|
||||
# the dataset from being finalized below.
|
||||
try:
|
||||
device.cleanup()
|
||||
except Exception:
|
||||
logging.exception("Device cleanup failed")
|
||||
try:
|
||||
if robot.is_connected:
|
||||
robot.disconnect()
|
||||
except Exception:
|
||||
logging.exception("Robot disconnect failed")
|
||||
|
||||
# Restore the terminal before the (potentially long) finalize/encode.
|
||||
if listener is not None:
|
||||
try:
|
||||
listener.stop()
|
||||
except Exception:
|
||||
logging.exception("Keyboard listener stop failed")
|
||||
|
||||
if dataset is not None:
|
||||
dataset.finalize()
|
||||
|
||||
if cfg.dataset.push_to_hub:
|
||||
if dataset is not None and dataset.num_episodes > 0:
|
||||
dataset.push_to_hub(tags=cfg.dataset.tags, private=cfg.dataset.private)
|
||||
else:
|
||||
logging.warning("No episodes saved — skipping push to hub")
|
||||
|
||||
logging.info("Exiting")
|
||||
|
||||
return dataset
|
||||
|
||||
|
||||
def main():
|
||||
record()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,117 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 NVIDIA Corporation and The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Teleoperate an SO-101 follower arm via NVIDIA Isaac Teleop.
|
||||
|
||||
``lerobot-teleoperate``-style CLI (draccus): ``--teleop.type`` selects the Isaac device
|
||||
(``xr_controller`` | ``so101_leader``), ``--robot.*`` the follower::
|
||||
|
||||
# XR (VR) controller: clutch + soft-orientation IK
|
||||
python -m examples.isaac_teleop_to_so101.teleoperate --robot.type=so101_follower \
|
||||
--robot.port=/dev/ttyACM0 --robot.id=so101_follower_arm --teleop.type=xr_controller
|
||||
|
||||
# SO-101 leader arm: 1:1 joint mirror (real leader on /dev/ttyACM1)
|
||||
python -m examples.isaac_teleop_to_so101.teleoperate --robot.type=so101_follower \
|
||||
--robot.port=/dev/ttyACM0 --robot.id=so101_follower_arm --teleop.type=so101_leader \
|
||||
--teleop.port=/dev/ttyACM1 --teleop.id=so101_leader_arm \
|
||||
--launch_plugin=/code/Teleop/install/plugins/so101_leader/so101_leader_plugin
|
||||
|
||||
``--teleop.type`` resolves against the Isaac device registry (see :class:`IsaacTeleopConfig`),
|
||||
distinct from the serial ``so101_leader``. The pipelines, clutch/IK/align internals, and
|
||||
reset-pose behavior live in ``common.py``. Requires the ``isaacteleop`` package and an OpenXR
|
||||
runtime (install instructions in this folder's ``README.md``).
|
||||
"""
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
from lerobot.configs import parser
|
||||
from lerobot.robots import RobotConfig
|
||||
from lerobot.robots.so_follower import SOFollowerConfig # noqa: F401 (registers so101_follower)
|
||||
from lerobot.utils.robot_utils import precise_sleep
|
||||
|
||||
from .common import (
|
||||
ALIGN_DURATION_S,
|
||||
FPS,
|
||||
RESET_DURATION_S,
|
||||
HoldLatch,
|
||||
build_device,
|
||||
)
|
||||
from .isaac_teleop import IsaacTeleopConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class TeleoperateConfig:
|
||||
"""``lerobot-teleoperate``-style CLI for the Isaac Teleop -> SO-101 example.
|
||||
|
||||
The fields below are the loop/launch knobs (not part of either device's config); the
|
||||
``[xr]`` / ``[leader]`` tags mark which device a knob applies to. Use ``--flag=false``
|
||||
for booleans (draccus style).
|
||||
"""
|
||||
|
||||
# Isaac Teleop input device + its knobs (--teleop.type=xr_controller|so101_leader,
|
||||
# then --teleop.<field>=...). Resolved against IsaacTeleopConfig's own choice registry.
|
||||
teleop: IsaacTeleopConfig
|
||||
# SO-101 FOLLOWER arm (--robot.type=so101_follower --robot.port=/dev/ttyACM0 --robot.id=...).
|
||||
robot: RobotConfig
|
||||
|
||||
# [leader] Path to the so101_leader plugin binary to spawn AFTER CloudXR is up (it then
|
||||
# inherits the runtime env). None (default) -> assume the plugin already runs externally.
|
||||
# The leader's serial port is --teleop.port (forwarded to the plugin; empty -> synthetic).
|
||||
launch_plugin: str | None = None
|
||||
|
||||
# [xr] Slew all joints to a default reset pose before the loop (--reset_to_origin=false to
|
||||
# keep the arm where it is). After the slew the clutch seeds its home from the measured pose.
|
||||
reset_to_origin: bool = True
|
||||
# [xr] Duration [s] of the reset-to-origin slew.
|
||||
reset_duration: float = RESET_DURATION_S
|
||||
|
||||
# [leader] Slew the follower to the leader's first pose before mirroring (--align=false to
|
||||
# begin the 1:1 mirror immediately; the follower may snap).
|
||||
align: bool = True
|
||||
# [leader] Duration [s] of the startup alignment slew.
|
||||
align_duration: float = ALIGN_DURATION_S
|
||||
|
||||
|
||||
@parser.wrap()
|
||||
def teleoperate(cfg: TeleoperateConfig):
|
||||
robot, device, motor_names = build_device(cfg)
|
||||
hold = HoldLatch(motor_names)
|
||||
try:
|
||||
while True:
|
||||
t0 = time.perf_counter()
|
||||
obs = robot.get_observation()
|
||||
# Idle (compute() -> None) holds the pose latched on the active->idle edge.
|
||||
action = hold.resolve(device.compute(obs), obs)
|
||||
robot.send_action(action)
|
||||
precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
finally:
|
||||
# A failing device cleanup must not skip the follower disconnect (which is what
|
||||
# disables torque on the arm).
|
||||
try:
|
||||
device.cleanup()
|
||||
finally:
|
||||
robot.disconnect()
|
||||
|
||||
|
||||
def main():
|
||||
teleoperate()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,79 @@
|
||||
#!/usr/bin/env python
|
||||
"""Convert a legacy LeRobot checkpoint to the current processor-pipeline format.
|
||||
|
||||
Older hub checkpoints (e.g. ``lerobot/act_aloha_sim_insertion_human``) bake
|
||||
normalization stats into the model weights and do not ship
|
||||
``policy_preprocessor.json`` / ``policy_postprocessor.json``. Current ``main``
|
||||
loads those processor configs from the checkpoint, so eval/rollout fail with
|
||||
``FileNotFoundError: Could not find 'policy_preprocessor.json'``.
|
||||
|
||||
This script rebuilds the processors from the training dataset's stats and saves
|
||||
a pipeline-format checkpoint locally that ``lerobot-eval`` can consume directly.
|
||||
|
||||
Usage:
|
||||
python examples/onnx/convert_legacy_checkpoint.py \
|
||||
--policy-path=lerobot/act_aloha_sim_insertion_human \
|
||||
--dataset-repo-id=lerobot/aloha_sim_insertion_human \
|
||||
--output-dir=outputs/converted/act_aloha_sim_insertion_human
|
||||
|
||||
Then:
|
||||
lerobot-eval \
|
||||
--policy.path=outputs/converted/act_aloha_sim_insertion_human \
|
||||
--env.type=aloha --env.task=AlohaInsertion-v0 \
|
||||
--eval.batch_size=10 --eval.n_episodes=50 \
|
||||
--eval.use_async_envs=false --policy.device=cuda
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
from lerobot.configs.policies import PreTrainedConfig
|
||||
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
||||
from lerobot.policies.factory import make_policy, make_pre_post_processors
|
||||
from lerobot.utils.constants import (
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--policy-path", required=True, help="Legacy checkpoint repo id or local dir")
|
||||
parser.add_argument(
|
||||
"--dataset-repo-id",
|
||||
required=True,
|
||||
help="Training dataset repo id, used only for normalization stats",
|
||||
)
|
||||
parser.add_argument("--output-dir", required=True, help="Where to save the converted checkpoint")
|
||||
parser.add_argument("--device", default="cpu", help="Device for building the policy (cpu is fine)")
|
||||
args = parser.parse_args()
|
||||
|
||||
out = Path(args.output_dir)
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"[1/4] Loading dataset stats from '{args.dataset_repo_id}' (metadata only)...")
|
||||
ds_meta = LeRobotDatasetMetadata(args.dataset_repo_id)
|
||||
|
||||
print(f"[2/4] Loading policy weights from '{args.policy_path}'...")
|
||||
cfg = PreTrainedConfig.from_pretrained(args.policy_path)
|
||||
cfg.pretrained_path = args.policy_path
|
||||
cfg.device = args.device
|
||||
policy = make_policy(cfg, ds_meta=ds_meta)
|
||||
|
||||
print("[3/4] Building processors from dataset stats...")
|
||||
preprocessor, postprocessor = make_pre_post_processors(
|
||||
policy_cfg=policy.config,
|
||||
dataset_stats=ds_meta.stats,
|
||||
)
|
||||
|
||||
print(f"[4/4] Saving pipeline-format checkpoint to '{out}'...")
|
||||
policy.save_pretrained(out)
|
||||
preprocessor.save_pretrained(out, config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json")
|
||||
postprocessor.save_pretrained(out, config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json")
|
||||
|
||||
print(f"\nDone. Converted checkpoint at: {out}")
|
||||
print("Eval it with --policy.path=" + str(out))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,178 @@
|
||||
#!/usr/bin/env python
|
||||
"""Evaluate an ACT policy in sim with either the PyTorch or ONNX network.
|
||||
|
||||
The ONNX backend swaps only ``policy.model`` (ResNet + transformer + action head)
|
||||
with an onnxruntime session. Everything else - the LeRobot processor pipeline
|
||||
(normalization), the action queue, and the gym env - is identical, so any
|
||||
difference in success rate is attributable to the network backend alone.
|
||||
|
||||
Run both backends with the same seed to compare:
|
||||
|
||||
python examples/onnx/eval_act_onnx.py \
|
||||
--policy-path=lerobot/act_aloha_sim_transfer_cube_human \
|
||||
--task=AlohaTransferCube-v0 \
|
||||
--backend=torch --n-episodes=50 --batch-size=10 --device=cuda
|
||||
|
||||
python examples/onnx/eval_act_onnx.py \
|
||||
--policy-path=lerobot/act_aloha_sim_transfer_cube_human \
|
||||
--task=AlohaTransferCube-v0 \
|
||||
--onnx=outputs/onnx/act_transfer_cube.onnx \
|
||||
--backend=onnx --n-episodes=50 --batch-size=10 --device=cuda
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from lerobot.envs.factory import make_env, make_env_config, make_env_pre_post_processors
|
||||
from lerobot.policies.act.modeling_act import ACTPolicy
|
||||
from lerobot.policies.factory import make_pre_post_processors
|
||||
from lerobot.scripts.lerobot_eval import eval_policy
|
||||
from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
|
||||
from lerobot.utils.random_utils import set_seed
|
||||
|
||||
|
||||
class ONNXACTModel(nn.Module):
|
||||
"""Drop-in replacement for ``ACTPolicy.model`` backed by onnxruntime."""
|
||||
|
||||
def __init__(
|
||||
self, onnx_path: str, image_keys: list[str], has_state: bool, has_env_state: bool, device: str
|
||||
):
|
||||
super().__init__()
|
||||
import onnxruntime as ort
|
||||
|
||||
providers = (
|
||||
["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
if str(device).startswith("cuda")
|
||||
else ["CPUExecutionProvider"]
|
||||
)
|
||||
so = ort.SessionOptions()
|
||||
so.log_severity_level = 3
|
||||
self.sess = ort.InferenceSession(onnx_path, sess_options=so, providers=providers)
|
||||
self.image_keys = image_keys
|
||||
self.has_state = has_state
|
||||
self.has_env_state = has_env_state
|
||||
print(f"[onnx] providers in use: {self.sess.get_providers()}")
|
||||
|
||||
def forward(self, batch: dict):
|
||||
state = batch[OBS_STATE] if self.has_state else batch[OBS_ENV_STATE]
|
||||
ref = state
|
||||
ort_inputs = {"state": state.detach().cpu().numpy().astype(np.float32)}
|
||||
images = batch[OBS_IMAGES]
|
||||
for i, img in enumerate(images):
|
||||
ort_inputs[f"image_{i}"] = img.detach().cpu().numpy().astype(np.float32)
|
||||
out = self.sess.run(None, ort_inputs)[0]
|
||||
actions = torch.from_numpy(out).to(ref.device, dtype=ref.dtype)
|
||||
return actions, None
|
||||
|
||||
|
||||
def load_stats_from_checkpoint(policy_path: str, input_features, output_features) -> dict:
|
||||
"""Recover MEAN_STD stats baked into a legacy ACT checkpoint's safetensors buffers.
|
||||
|
||||
Legacy checkpoints store normalization as buffers like
|
||||
``normalize_inputs.buffer_observation_state.{mean,std}``. We map those back to
|
||||
feature names so we can rebuild the processor pipeline without the dataset.
|
||||
"""
|
||||
from safetensors.torch import load_file
|
||||
|
||||
p = Path(policy_path)
|
||||
if p.is_dir():
|
||||
st_path = p / "model.safetensors"
|
||||
else:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
st_path = Path(hf_hub_download(policy_path, "model.safetensors"))
|
||||
|
||||
sd = load_file(str(st_path))
|
||||
stats: dict = {}
|
||||
for feat in list(input_features) + list(output_features):
|
||||
buf = "buffer_" + feat.replace(".", "_")
|
||||
for prefix in ("normalize_inputs", "normalize_targets", "unnormalize_outputs"):
|
||||
mkey, skey = f"{prefix}.{buf}.mean", f"{prefix}.{buf}.std"
|
||||
if mkey in sd and skey in sd:
|
||||
stats[feat] = {"mean": sd[mkey].numpy(), "std": sd[skey].numpy()}
|
||||
break
|
||||
return stats
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--policy-path", required=True)
|
||||
parser.add_argument("--task", required=True, help="e.g. AlohaTransferCube-v0")
|
||||
parser.add_argument("--env-type", default="aloha")
|
||||
parser.add_argument("--backend", choices=["torch", "onnx"], default="torch")
|
||||
parser.add_argument("--onnx", default=None, help="Path to .onnx (required for --backend=onnx)")
|
||||
parser.add_argument("--n-episodes", type=int, default=50)
|
||||
parser.add_argument("--batch-size", type=int, default=10)
|
||||
parser.add_argument("--device", default="cuda")
|
||||
parser.add_argument("--seed", type=int, default=1000)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.backend == "onnx" and not args.onnx:
|
||||
raise SystemExit("--backend=onnx requires --onnx=<path>")
|
||||
|
||||
device = "cuda" if (args.device == "cuda" and torch.cuda.is_available()) else "cpu"
|
||||
set_seed(args.seed)
|
||||
|
||||
print(f"[1/4] Loading ACT policy from '{args.policy_path}'...")
|
||||
policy = ACTPolicy.from_pretrained(args.policy_path)
|
||||
policy.config.device = device
|
||||
policy.eval()
|
||||
policy.to(device)
|
||||
cfg = policy.config
|
||||
|
||||
if args.backend == "onnx":
|
||||
image_keys = list(cfg.image_features)
|
||||
has_state = cfg.robot_state_feature is not None
|
||||
has_env_state = cfg.env_state_feature is not None
|
||||
print(f"[2/4] Swapping policy.model with ONNX backend ({args.onnx})")
|
||||
policy.model = ONNXACTModel(args.onnx, image_keys, has_state, has_env_state, device)
|
||||
policy.to(device)
|
||||
else:
|
||||
print("[2/4] Using PyTorch backend")
|
||||
|
||||
print("[3/4] Building processors and environment...")
|
||||
stats = load_stats_from_checkpoint(args.policy_path, cfg.input_features, cfg.output_features)
|
||||
preprocessor, postprocessor = make_pre_post_processors(
|
||||
policy_cfg=cfg,
|
||||
dataset_stats=stats,
|
||||
preprocessor_overrides={"device_processor": {"device": device}},
|
||||
)
|
||||
|
||||
env_cfg = make_env_config(args.env_type, task=args.task)
|
||||
env_preprocessor, env_postprocessor = make_env_pre_post_processors(env_cfg=env_cfg, policy_cfg=cfg)
|
||||
env_groups = make_env(env_cfg, n_envs=args.batch_size, use_async_envs=False)
|
||||
# make_env returns {task_group: {idx: VectorEnv}}; grab the single env.
|
||||
first_group = next(iter(env_groups.values()))
|
||||
env = next(iter(first_group.values()))
|
||||
|
||||
print(f"[4/4] Evaluating backend='{args.backend}' for {args.n_episodes} episodes (seed={args.seed})...")
|
||||
with torch.no_grad():
|
||||
info = eval_policy(
|
||||
env=env,
|
||||
policy=policy,
|
||||
env_preprocessor=env_preprocessor,
|
||||
env_postprocessor=env_postprocessor,
|
||||
preprocessor=preprocessor,
|
||||
postprocessor=postprocessor,
|
||||
n_episodes=args.n_episodes,
|
||||
start_seed=args.seed,
|
||||
)
|
||||
|
||||
agg = info["aggregated"]
|
||||
print("\n==== RESULT ====")
|
||||
print(f"backend : {args.backend}")
|
||||
print(f"task : {args.task}")
|
||||
print(f"n_episodes : {args.n_episodes}")
|
||||
print(f"pc_success : {agg['pc_success']:.1f}%")
|
||||
print(f"avg_max_reward: {agg['avg_max_reward']:.4f}")
|
||||
print(f"eval_ep_s : {agg['eval_ep_s']:.2f}s")
|
||||
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,133 @@
|
||||
#!/usr/bin/env python
|
||||
"""Export an ACT policy's network to ONNX and verify numerical parity.
|
||||
|
||||
Only the inference network is exported (ResNet backbone + transformer enc/dec +
|
||||
action head). The VAE encoder is training-only and the inference latent is zeros,
|
||||
so the exported graph is a pure function of (state, images) -> action_chunk.
|
||||
Normalization stays in the LeRobot processor pipeline (outside ONNX).
|
||||
|
||||
Usage:
|
||||
python examples/onnx/export_act.py \
|
||||
--policy-path=outputs/converted/act_aloha_sim_transfer_cube_human \
|
||||
--output=outputs/onnx/act_transfer_cube.onnx
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from lerobot.policies.act.modeling_act import ACTPolicy
|
||||
from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
|
||||
|
||||
|
||||
class ACTExportWrapper(nn.Module):
|
||||
"""Tensor-in/tensor-out wrapper around ACT's inference network."""
|
||||
|
||||
def __init__(self, model: nn.Module, image_keys: list[str], has_state: bool, has_env_state: bool):
|
||||
super().__init__()
|
||||
self.model = model
|
||||
self.image_keys = image_keys
|
||||
self.has_state = has_state
|
||||
self.has_env_state = has_env_state
|
||||
|
||||
def forward(self, state: torch.Tensor, *images: torch.Tensor) -> torch.Tensor:
|
||||
batch: dict = {}
|
||||
if self.has_state:
|
||||
batch[OBS_STATE] = state
|
||||
if self.has_env_state:
|
||||
# Convention: when env_state is used it is passed as `state`.
|
||||
batch[OBS_ENV_STATE] = state
|
||||
batch[OBS_IMAGES] = list(images)
|
||||
actions, _ = self.model(batch)
|
||||
return actions
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--policy-path", required=True, help="Converted ACT checkpoint dir or repo id")
|
||||
parser.add_argument("--output", required=True, help="Output .onnx path")
|
||||
parser.add_argument("--opset", type=int, default=17)
|
||||
parser.add_argument("--atol", type=float, default=1e-3)
|
||||
parser.add_argument("--device", default="cpu")
|
||||
args = parser.parse_args()
|
||||
|
||||
out = Path(args.output)
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"[1/4] Loading ACT policy from '{args.policy_path}'...")
|
||||
policy = ACTPolicy.from_pretrained(args.policy_path)
|
||||
policy.eval()
|
||||
policy.to(args.device)
|
||||
cfg = policy.config
|
||||
|
||||
image_keys = list(cfg.image_features)
|
||||
has_state = cfg.robot_state_feature is not None
|
||||
has_env_state = cfg.env_state_feature is not None
|
||||
state_dim = (cfg.robot_state_feature or cfg.env_state_feature).shape[0]
|
||||
|
||||
print(
|
||||
f" image_keys={image_keys} state_dim={state_dim} "
|
||||
f"chunk_size={cfg.chunk_size} action_dim={cfg.action_feature.shape[0]}"
|
||||
)
|
||||
|
||||
wrapper = ACTExportWrapper(policy.model, image_keys, has_state, has_env_state).eval().to(args.device)
|
||||
|
||||
# Build example inputs (batch size 1) from the config feature shapes.
|
||||
state_example = torch.randn(1, state_dim, device=args.device)
|
||||
image_examples = [torch.rand(1, *cfg.image_features[k].shape, device=args.device) for k in image_keys]
|
||||
example_inputs = (state_example, *image_examples)
|
||||
|
||||
input_names = ["state"] + [f"image_{i}" for i in range(len(image_keys))]
|
||||
output_names = ["action_chunk"]
|
||||
dynamic_axes = {name: {0: "batch"} for name in input_names + output_names}
|
||||
|
||||
print(f"[2/4] Exporting to ONNX (opset {args.opset}) -> {out}")
|
||||
torch.onnx.export(
|
||||
wrapper,
|
||||
example_inputs,
|
||||
str(out),
|
||||
input_names=input_names,
|
||||
output_names=output_names,
|
||||
dynamic_axes=dynamic_axes,
|
||||
opset_version=args.opset,
|
||||
do_constant_folding=True,
|
||||
dynamo=False,
|
||||
)
|
||||
|
||||
print("[3/4] Running parity check (torch vs onnxruntime)...")
|
||||
import onnxruntime as ort
|
||||
|
||||
providers = ["CPUExecutionProvider"]
|
||||
so = ort.SessionOptions()
|
||||
so.log_severity_level = 3
|
||||
sess = ort.InferenceSession(str(out), sess_options=so, providers=providers)
|
||||
|
||||
# Fresh random inputs for the check.
|
||||
state_check = torch.randn(2, state_dim, device=args.device)
|
||||
image_check = [torch.rand(2, *cfg.image_features[k].shape, device=args.device) for k in image_keys]
|
||||
|
||||
with torch.no_grad():
|
||||
torch_out = wrapper(state_check, *image_check).cpu().numpy()
|
||||
|
||||
ort_inputs = {"state": state_check.cpu().numpy()}
|
||||
for i, img in enumerate(image_check):
|
||||
ort_inputs[f"image_{i}"] = img.cpu().numpy()
|
||||
ort_out = sess.run(None, ort_inputs)[0]
|
||||
|
||||
max_abs = float(np.max(np.abs(torch_out - ort_out)))
|
||||
mean_abs = float(np.mean(np.abs(torch_out - ort_out)))
|
||||
print(f" shapes: torch={torch_out.shape} onnx={ort_out.shape}")
|
||||
print(f" max_abs_diff={max_abs:.3e} mean_abs_diff={mean_abs:.3e} (atol={args.atol:.0e})")
|
||||
|
||||
ok = max_abs <= args.atol
|
||||
print(f"[4/4] Parity: {'PASS' if ok else 'FAIL'}")
|
||||
if not ok:
|
||||
raise SystemExit(f"Parity check failed: max_abs_diff {max_abs:.3e} > atol {args.atol:.0e}")
|
||||
print(f"\nDone. ONNX model at: {out}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+1
-1
@@ -25,7 +25,7 @@ discord = "https://discord.gg/s3KuuzsPFb"
|
||||
|
||||
[project]
|
||||
name = "lerobot"
|
||||
version = "0.6.1"
|
||||
version = "0.5.2"
|
||||
description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch"
|
||||
dynamic = ["readme"]
|
||||
license = { text = "Apache-2.0" }
|
||||
|
||||
@@ -65,14 +65,6 @@ class PlanConfig:
|
||||
# invented from the task text (+1 VLM call/episode).
|
||||
subtask_describe_first: bool = True
|
||||
|
||||
# Seeded relabeling: after segmentation, re-label each span with a focused
|
||||
# pass that sees the previous / current / next segment contact sheets and
|
||||
# minimally corrects the seed label (macrodata's best end-to-end labeling
|
||||
# step). Costs +1 VLM call per subtask; off by default.
|
||||
subtask_seeded_relabel: bool = False
|
||||
# Frames sampled uniformly per segment sheet in the relabel pass.
|
||||
subtask_relabel_frames: int = 5
|
||||
|
||||
# Emit ``style="plan"`` rows at each boundary; False = subtasks + memory only.
|
||||
emit_plan: bool = True
|
||||
|
||||
@@ -168,11 +160,6 @@ class VlmConfig:
|
||||
# Forwarded as extra_body.chat_template_kwargs (e.g. {"enable_thinking": false}).
|
||||
chat_template_kwargs: dict[str, Any] | None = None
|
||||
|
||||
# OpenAI-style thinking budget hint ("low"/"medium"/"high"); forwarded to
|
||||
# the server when set. Used to cap a thinking model's reasoning so it
|
||||
# leaves tokens for the actual JSON answer on OpenAI-compatible endpoints.
|
||||
reasoning_effort: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExecutorConfig:
|
||||
|
||||
@@ -413,16 +413,7 @@ def _draw_timestamp_badge(image: PIL.Image.Image, timestamp: float) -> PIL.Image
|
||||
|
||||
result = image.copy()
|
||||
draw = ImageDraw.Draw(result)
|
||||
# Scale the timestamp to the tile so it stays legible after the model
|
||||
# downsamples the full sheet into 768px tiles — a tiny bitmap font blurs
|
||||
# at contact-sheet resolution and the VLM can no longer read the exact
|
||||
# source time, which is what the boundary score depends on. ``size=`` is
|
||||
# supported by Pillow's bitmap default since 10.1; fall back otherwise.
|
||||
badge_px = max(14, round(image.height * 0.12))
|
||||
try:
|
||||
font = ImageFont.load_default(size=badge_px)
|
||||
except TypeError:
|
||||
font = ImageFont.load_default()
|
||||
font = ImageFont.load_default()
|
||||
label = f"{timestamp:06.2f}s"
|
||||
left, top, right, bottom = draw.textbbox((0, 0), label, font=font)
|
||||
text_w, text_h = right - left, bottom - top
|
||||
|
||||
@@ -116,8 +116,6 @@ class PlanSubtasksMemoryModule:
|
||||
rows.extend(self._task_aug_rows([effective_task, *variants], t0))
|
||||
|
||||
subtask_spans = self._generate_subtasks(record, task=effective_task)
|
||||
if self.config.subtask_seeded_relabel and subtask_spans:
|
||||
subtask_spans = self._seeded_relabel(record, subtask_spans, effective_task)
|
||||
|
||||
# subtask rows
|
||||
for span in subtask_spans:
|
||||
@@ -511,51 +509,6 @@ class PlanSubtasksMemoryModule:
|
||||
|
||||
return cleaned
|
||||
|
||||
def _seeded_relabel(
|
||||
self, record: EpisodeRecord, spans: list[dict[str, Any]], task: str
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Re-label each span using prev/current/next segment contact sheets.
|
||||
|
||||
Boundaries are kept fixed; only ``text`` is refined. The original
|
||||
("seed") label is passed as a strong prior so the model verifies and
|
||||
minimally corrects it rather than re-describing from scratch — the
|
||||
macrodata seeded-relabeling step. One VLM call per span.
|
||||
"""
|
||||
n = len(spans)
|
||||
out: list[dict[str, Any]] = []
|
||||
for i, span in enumerate(spans):
|
||||
content: list[dict[str, Any]] = []
|
||||
if i > 0:
|
||||
content += self._segment_sheet(record, spans[i - 1])
|
||||
content += self._segment_sheet(record, span)
|
||||
if i < n - 1:
|
||||
content += self._segment_sheet(record, spans[i + 1])
|
||||
prompt = load_prompt("plan_subtask_relabel").format(
|
||||
episode_task=task,
|
||||
seed_label=span["text"],
|
||||
segment_index=i + 1,
|
||||
segment_count=n,
|
||||
start=float(span["start"]),
|
||||
end=float(span["end"]),
|
||||
)
|
||||
content.append({"type": "text", "text": prompt})
|
||||
label = self._vlm_field([{"role": "user", "content": content}], "label")
|
||||
text = label.strip() if isinstance(label, str) and label.strip() else span["text"]
|
||||
out.append({**span, "text": text})
|
||||
return out
|
||||
|
||||
def _segment_sheet(self, record: EpisodeRecord, span: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""Contact-sheet block(s) for one span: up to N frames sampled uniformly."""
|
||||
s, e = float(span["start"]), float(span["end"])
|
||||
n = max(1, int(self.config.subtask_relabel_frames))
|
||||
if e <= s or n == 1:
|
||||
timestamps = [s]
|
||||
else:
|
||||
step = (e - s) / (n - 1)
|
||||
timestamps = [s + i * step for i in range(n)]
|
||||
frames = self.frame_provider.frames_at(record, timestamps)
|
||||
return self._contact_sheet_blocks(frames, timestamps[: len(frames)])
|
||||
|
||||
def _generate_subtasks_windowed(
|
||||
self, record: EpisodeRecord, task: str, window_s: float
|
||||
) -> list[dict[str, Any]]:
|
||||
|
||||
@@ -22,23 +22,12 @@ plain editors and roundtrip cleanly through ``ruff format``.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
_DIR = Path(__file__).parent
|
||||
|
||||
|
||||
def load(name: str) -> str:
|
||||
"""Read prompt template ``name.txt`` from the ``prompts/`` directory.
|
||||
|
||||
A ``LEROBOT_PROMPT_OVERRIDE_<name>`` environment variable, when set to a
|
||||
non-empty value, takes precedence over the packaged file. This lets prompt
|
||||
search (e.g. GEPA) inject candidate templates into a remote job without
|
||||
rebuilding the package; the override must keep the same ``{placeholder}``
|
||||
fields the call site formats in.
|
||||
"""
|
||||
override = os.environ.get(f"LEROBOT_PROMPT_OVERRIDE_{name}")
|
||||
if override and override.strip():
|
||||
return override
|
||||
"""Read prompt template ``name.txt`` from the ``prompts/`` directory."""
|
||||
path = _DIR / f"{name}.txt"
|
||||
return path.read_text(encoding="utf-8")
|
||||
|
||||
@@ -1,35 +0,0 @@
|
||||
Annotate one fixed segment from a longer robot demonstration.
|
||||
|
||||
Return only JSON:
|
||||
{{"label": "<short descriptive subtask label>"}}
|
||||
|
||||
You are shown up to three timestamped contact sheets, in order:
|
||||
- The FIRST sheet is the PREVIOUS segment (context only); it may be absent.
|
||||
- The SECOND sheet is the CURRENT target segment.
|
||||
- The THIRD sheet is the NEXT segment (context only); it may be absent.
|
||||
Each tile has its timestamp (seconds, absolute video time) burned into its
|
||||
top-left corner.
|
||||
|
||||
Episode instruction: "{episode_task}"
|
||||
Target segment: {segment_index} of {segment_count}
|
||||
Target time: {start:.2f}s to {end:.2f}s
|
||||
Original predicted label for this exact segment: "{seed_label}"
|
||||
|
||||
Rules:
|
||||
- Label ONLY the current target segment (the second sheet). Use the
|
||||
previous/next sheets only to disambiguate what changed.
|
||||
- Treat the original predicted label as a STRONG PRIOR, not ground truth:
|
||||
verify it against the current segment and correct it minimally.
|
||||
- If it already names the right action and main object, keep it; only fix
|
||||
grammar or add a clearly visible essential detail.
|
||||
- If it is vague but directionally correct, make it more specific.
|
||||
- If it describes the previous/next segment, the wrong action, wrong
|
||||
object, wrong destination, or a wrong state change, replace it.
|
||||
- Do not describe the previous or next segment, and do not split, merge,
|
||||
or move the fixed segment.
|
||||
- Do not introduce an action that is not clearly visible in the current
|
||||
target segment.
|
||||
- Use one concise imperative phrase. Name the manipulated object and the
|
||||
action / state change. Include source, destination, side, direction,
|
||||
final placement, or opened/closed state when visible and central.
|
||||
- Do not mention timestamps, frame numbers, uncertainty, or intent.
|
||||
@@ -1,68 +1,112 @@
|
||||
You are annotating a teleoperated robot demonstration shown as
|
||||
timestamped contact sheets (each tile has its time in seconds burned
|
||||
into the top-left corner). The operator's goal was: "{episode_task}"
|
||||
You are labeling a teleoperated robot demonstration.
|
||||
|
||||
{observation_block}Reconstruct the sequence of COMPLETED manipulation events the robot
|
||||
performs, in chronological order. Output one segment per event with a
|
||||
[start, end] time in seconds and a short action label.
|
||||
The user originally asked: "{episode_task}"
|
||||
|
||||
GROUNDING — read first, it overrides everything below:
|
||||
- Label ONLY events you can SEE in the frames. The instruction is the
|
||||
goal; the VIDEO is the ground truth for what actually happened.
|
||||
- Do NOT invent, anticipate, or pad steps that are not shown.
|
||||
You are shown the entire demonstration as a single video. Watch the
|
||||
whole clip, then segment it into a list of consecutive atomic subtasks
|
||||
the robot performs.
|
||||
|
||||
Granularity — segment by completed events, not by motion:
|
||||
- Start a NEW segment whenever the world state changes: an object is
|
||||
grasped, lifted, transported, placed, or released; a held object
|
||||
changes; a drawer/door/lid/container opens or closes; contents move
|
||||
between containers (poured); a tool starts or stops acting on a
|
||||
surface. Watch the gripper open/close transitions — they usually mark
|
||||
boundaries.
|
||||
- Do NOT split approach, reach, grasp adjustment, small repositioning,
|
||||
hesitation, or retreat into their own segments. Fold each into the
|
||||
event it belongs to (the approach is part of the pick; the retreat is
|
||||
part of the place).
|
||||
- Do NOT merge separate completed events. Each distinct pick, place,
|
||||
open, close, pour, push, wipe, or insert is its own segment, even when
|
||||
they repeat on different objects or locations.
|
||||
- Most segments last 2-10 seconds. Shorter segments are okay ONLY for
|
||||
fast pick / place / open / close / release events. Never emit a
|
||||
segment shorter than {min_subtask_seconds} seconds; merge a too-short
|
||||
candidate into its neighbour instead.
|
||||
- Skip idle time, pure camera motion, and tiny hand jitter.
|
||||
{observation_block}GROUNDING — read this first, it overrides everything below:
|
||||
- Label ONLY what the robot actually does in the video. Every subtask
|
||||
you emit must correspond to motion you can SEE in specific frames.
|
||||
- Do NOT invent, anticipate, or pad. If the robot only does one thing
|
||||
(e.g. it just navigates to a location and the clip ends), emit
|
||||
EXACTLY ONE subtask. Many demonstrations are a single atomic skill.
|
||||
- ``max_steps`` below is a hard CEILING, not a target. Emitting fewer
|
||||
subtasks than the ceiling is not just allowed, it is expected for
|
||||
short / atomic demonstrations. One correct subtask is far better
|
||||
than several invented ones.
|
||||
- If the video does not clearly show the action implied by the task,
|
||||
describe what you actually see — do NOT fabricate the task's steps
|
||||
from the instruction text. The instruction tells you the goal; the
|
||||
VIDEO is the ground truth for what happened.
|
||||
|
||||
Labels — short imperative phrases:
|
||||
- One concise command naming the action and the manipulated object, e.g.
|
||||
"pick up the red cup", "put the cup on the shelf", "open the top
|
||||
drawer", "pour water into the glass", "insert the plug into the
|
||||
socket".
|
||||
- Include source, destination, side, direction, or the final
|
||||
open/closed state when it is visible and central to the event.
|
||||
- Prefer these verbs (extend only when none fits): pick up, put, place,
|
||||
push, pull, turn, press, open, close, pour, insert, wipe, stack.
|
||||
Disambiguate by what you SEE:
|
||||
* STACK vs PUT: object placed ON TOP OF another object -> "stack".
|
||||
* INSERT vs PUT: object pushed INTO a fitted slot/hole/socket -> "insert".
|
||||
* PICK UP vs PUT (direction): gripper CLOSES and object moves WITH
|
||||
the hand -> "pick up"; gripper OPENS and object stays -> "put".
|
||||
* POUR vs PUT: source is tilted and contents flow -> "pour".
|
||||
- Use the exact object nouns implied by the task; stay consistent across
|
||||
the episode (don't switch "cube" to "block").
|
||||
- Write imperative commands, never third person ("the robot ..."), and
|
||||
drop articles/adverbs.
|
||||
Authoring rules — Hi Robot atom granularity, pi0.7-style short prompts:
|
||||
|
||||
Timing:
|
||||
- Use the burned-in timestamps to set start and end. Boundaries should
|
||||
land on or near a printed time, and every [start, end] must lie within
|
||||
[0.0, {episode_duration}] seconds, be non-overlapping, and cover the
|
||||
episode in order.
|
||||
- Emit at most {max_steps} segments.
|
||||
- Each subtask = one COMPOSITE atomic skill the low-level policy can
|
||||
execute end-to-end. A "skill" bundles its own approach motion with
|
||||
its terminal action — do NOT split the approach off as its own
|
||||
subtask. The whole-arm policy already learns to reach as part of
|
||||
every manipulation primitive.
|
||||
- Write each subtask as an IMPERATIVE COMMAND, starting with one of
|
||||
these verbs (extend only when none fits):
|
||||
pick up <obj> — approach + grasp + lift in one subtask
|
||||
put <obj> on/in <loc> — transport + release in one subtask
|
||||
place <obj> on/in <loc> — synonym of "put"; pick one and stay consistent
|
||||
push <obj> — contact + linear shove
|
||||
pull <obj> — contact + linear retract
|
||||
turn <knob/dial/handle> — rotary actuation
|
||||
press <button> — single-press contact
|
||||
open <drawer/door/lid> — full open motion
|
||||
close <drawer/door/lid> — full close motion
|
||||
pour <src> into <dst> — tilt + flow
|
||||
insert <obj> into <slot>— alignment + push-fit
|
||||
go to <loc> — ONLY when no grasp / actuation follows
|
||||
(e.g. a pure relocation between phases).
|
||||
If the next subtask grasps something at
|
||||
that location, drop "go to ..." and just
|
||||
write "pick up ..." instead.
|
||||
- Forbidden ultra-fine splits — the VLM is NOT allowed to emit these
|
||||
as standalone subtasks; fold them into the parent composite:
|
||||
"move to X" → fold into "pick up X" (or whatever follows)
|
||||
"reach for X" → fold into "pick up X"
|
||||
"grasp X" → fold into "pick up X"
|
||||
"lift X" → fold into "pick up X" (or "put X on Y" if it's
|
||||
the transport phase of a place)
|
||||
"release X" → fold into "put X on Y" (or "place X in Y")
|
||||
- Keep it SHORT — a verb phrase, not a sentence. Drop articles
|
||||
("the", "a") and adverbs ("carefully", "slowly"). Add a "how"
|
||||
detail (which hand, which grasp point) ONLY when it is needed to
|
||||
disambiguate. Every subtask must begin with one of the verbs
|
||||
above (no leading nouns, no "then", no "first").
|
||||
- NEVER use third person. Never write "the robot", "the arm", "the
|
||||
gripper moves", "it picks up" — the robot is implied. Command it,
|
||||
do not describe it.
|
||||
- Use the exact object nouns from the task above. If the task says
|
||||
"cube", every subtask says "cube" — never switch to "block". If it
|
||||
says "box", never switch to "bin"/"container". Keep vocabulary
|
||||
consistent across the whole episode.
|
||||
- Good: "pick up blue cube", "put blue cube in box", "open drawer",
|
||||
"turn red knob", "press start button", "go to sink".
|
||||
- Bad: "move to blue cube" (approach as its own subtask — forbidden,
|
||||
must be folded into "pick up blue cube"); "the robot arm moves
|
||||
towards the blue cube" (third person, too long); "carefully pick
|
||||
up the cube" (adverb, article); "release the yellow block"
|
||||
("block" when the task said "cube", and "release" must be folded
|
||||
into a "put"/"place" subtask).
|
||||
- Subtasks are non-overlapping and cover the full episode in order.
|
||||
Choose the cut points yourself based on what you see in the video
|
||||
(gripper open/close events, contact, regrasps, transitions).
|
||||
- Each subtask spans at least {min_subtask_seconds} seconds. If a
|
||||
candidate span would be shorter, merge it into its neighbour
|
||||
rather than emitting it.
|
||||
- Do not exceed {max_steps} subtasks total. Fewer, larger composites
|
||||
are preferred over many micro-steps.
|
||||
- Every subtask's [start_time, end_time] must lie within
|
||||
[0.0, {episode_duration}] seconds.
|
||||
|
||||
SPECIAL CASES — verb disambiguation (each rule is narrowly visual and
|
||||
fires ONLY on the spatial situation it names; it must not change how you
|
||||
label any other situation):
|
||||
- STACK vs PUT: if an object is placed ON TOP OF another specific object
|
||||
(not on a flat table / shelf / counter), use "stack ... on ...", not
|
||||
"put". "stack blue book on green book", NOT "put blue book on table".
|
||||
- INSERT vs PUT: if an object goes INTO a fitted slot / hole / socket /
|
||||
receptacle (push-fit), use "insert ... into ...", not "put".
|
||||
- RETRIEVE/PICK-UP vs PUT (direction): watch the gripper. If it CLOSES
|
||||
on the object and the object moves WITH the hand, it is "pick up" /
|
||||
"retrieve" (object leaves its location). If the gripper OPENS and the
|
||||
object stays where the hand left it, it is "put" / "place" (object
|
||||
arrives at a location). Decide by which way the object moves, not by
|
||||
where the hand ends up.
|
||||
- POUR vs PUT: only use "pour" when the source is tilted and contents
|
||||
flow out; moving a full container without tilting is "put"/"place".
|
||||
|
||||
Output strictly valid JSON of shape:
|
||||
|
||||
{{
|
||||
"subtasks": [
|
||||
{{"text": "<short imperative action label>", "start": <float>, "end": <float>}},
|
||||
{{"text": "<short imperative verb phrase>", "start": <float>, "end": <float>}},
|
||||
...
|
||||
]
|
||||
}}
|
||||
|
||||
@@ -285,8 +285,6 @@ def _make_openai_client(config: VlmConfig) -> VlmClient:
|
||||
"max_tokens": max_tok,
|
||||
"temperature": temp,
|
||||
}
|
||||
if config.reasoning_effort:
|
||||
kwargs["reasoning_effort"] = config.reasoning_effort
|
||||
extra_body: dict[str, Any] = {}
|
||||
if send_mm_kwargs and mm_kwargs:
|
||||
extra_body["mm_processor_kwargs"] = {**mm_kwargs, "do_sample_frames": True}
|
||||
@@ -298,13 +296,7 @@ def _make_openai_client(config: VlmConfig) -> VlmClient:
|
||||
chosen = clients[rr_counter["i"] % len(clients)]
|
||||
rr_counter["i"] += 1
|
||||
response = chosen.chat.completions.create(**kwargs)
|
||||
# Some OpenAI-compatible servers can return a choice with no message
|
||||
# (safety filter, or a "thinking" model that spends the whole budget
|
||||
# before emitting content). Treat that as an empty reply so the
|
||||
# JSON-retry path handles it instead of crashing the run.
|
||||
choice = response.choices[0] if response.choices else None
|
||||
message = choice.message if choice is not None else None
|
||||
return (message.content if message is not None else None) or ""
|
||||
return response.choices[0].message.content or ""
|
||||
|
||||
def _gen(batch: Sequence[Sequence[dict[str, Any]]], max_tok: int, temp: float) -> list[str]:
|
||||
if len(batch) <= 1 or config.client_concurrency <= 1:
|
||||
|
||||
@@ -205,30 +205,24 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
|
||||
f"{CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
|
||||
) from e
|
||||
|
||||
# HACK: Parse the original config to get the config subclass, so that we can
|
||||
# apply cli overrides.
|
||||
# This is very ugly, ideally we'd like to be able to do that natively with draccus
|
||||
# something like --policy.path (in addition to --policy.type)
|
||||
with draccus.config_type("json"):
|
||||
orig_config = draccus.parse(cls, config_file, args=[])
|
||||
|
||||
if config_file is None:
|
||||
raise FileNotFoundError(f"{CONFIG_NAME} not found in {model_id}")
|
||||
|
||||
with open(config_file) as f:
|
||||
config = json.load(f)
|
||||
|
||||
# Resolve the concrete config subclass from the serialized "type" tag, then parse
|
||||
# the config (with CLI overrides) directly for that class. The "type" key is
|
||||
# stripped because draccus only consumes it when parsing the registry base class.
|
||||
policy_type = config.pop("type", None)
|
||||
if policy_type is None:
|
||||
raise ValueError(f"Missing 'type' field in {CONFIG_NAME} of {model_id}")
|
||||
try:
|
||||
config_cls = cls.get_choice_class(policy_type)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Policy type '{policy_type}' (from {CONFIG_NAME} of {model_id}) is not registered. "
|
||||
f"Available policy types: {cls.get_known_choices()}"
|
||||
) from e
|
||||
|
||||
config.pop("type")
|
||||
with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as f:
|
||||
json.dump(config, f)
|
||||
config_file = f.name
|
||||
|
||||
cli_overrides = policy_kwargs.pop("cli_overrides", [])
|
||||
with draccus.config_type("json"):
|
||||
return draccus.parse(config_cls, config_file, args=cli_overrides)
|
||||
return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)
|
||||
|
||||
@@ -32,7 +32,6 @@ from .pretrained import PreTrainedPolicy as PreTrainedPolicy
|
||||
from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
|
||||
from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
|
||||
from .utils import make_robot_action, prepare_observation_for_inference
|
||||
from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig as VLAJEPAConfig
|
||||
from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
|
||||
from .wall_x.configuration_wall_x import WallXConfig as WallXConfig
|
||||
from .xvla.configuration_xvla import XVLAConfig as XVLAConfig
|
||||
@@ -58,7 +57,6 @@ __all__ = [
|
||||
"PI05Config",
|
||||
"SmolVLAConfig",
|
||||
"TDMPCConfig",
|
||||
"VLAJEPAConfig",
|
||||
"VQBeTConfig",
|
||||
"WallXConfig",
|
||||
"XVLAConfig",
|
||||
|
||||
@@ -18,10 +18,17 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
make_default_pre_post_processors,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_act import ACTConfig
|
||||
|
||||
@@ -47,4 +54,34 @@ def make_act_pre_post_processors(
|
||||
tuple[PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], PolicyProcessorPipeline[PolicyAction, PolicyAction]]: A tuple containing the
|
||||
pre-processor pipeline and the post-processor pipeline.
|
||||
"""
|
||||
return make_default_pre_post_processors(config, dataset_stats, normalizer_device=config.device)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
device=config.device,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -19,10 +19,17 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
make_default_pre_post_processors,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_diffusion import DiffusionConfig
|
||||
|
||||
@@ -56,4 +63,32 @@ def make_diffusion_pre_post_processors(
|
||||
Returns:
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
return make_default_pre_post_processors(config, dataset_stats)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -23,16 +23,24 @@ import torch
|
||||
|
||||
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
ComplementaryDataProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
)
|
||||
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
|
||||
from lerobot.types import TransitionKey
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
from lerobot.utils.constants import (
|
||||
OBS_STATE,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
from lerobot.utils.import_utils import _transformers_available, require_package
|
||||
|
||||
from .configuration_eo1 import EO1Config
|
||||
@@ -234,12 +242,14 @@ def make_eo1_pre_post_processors(
|
||||
]:
|
||||
"""Build pre/post processor pipelines for EO1."""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps: list[ProcessorStep] = [
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
steps.normalize,
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
EO1ConversationTemplateStep(input_features=config.input_features, chunk_size=config.chunk_size),
|
||||
EO1QwenProcessorStep(
|
||||
processor_name=config.vlm_base,
|
||||
@@ -247,12 +257,27 @@ def make_eo1_pre_post_processors(
|
||||
image_max_pixels=config.image_max_pixels,
|
||||
use_fast_processor=config.use_fast_processor,
|
||||
),
|
||||
steps.to_device,
|
||||
DeviceProcessorStep(device=config.device),
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
steps.unnormalize,
|
||||
steps.to_cpu,
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -27,11 +27,9 @@ from lerobot.utils.import_utils import _transformers_available, require_package
|
||||
|
||||
if TYPE_CHECKING or _transformers_available:
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
from transformers.utils import is_flash_attn_2_available
|
||||
else:
|
||||
AutoModel = None
|
||||
AutoTokenizer = None
|
||||
is_flash_attn_2_available = None
|
||||
|
||||
IMAGENET_MEAN = (0.485, 0.456, 0.406)
|
||||
IMAGENET_STD = (0.229, 0.224, 0.225)
|
||||
@@ -137,13 +135,9 @@ class InternVL3Embedder(nn.Module):
|
||||
raise ValueError(f"Unsupported EVO1 vlm_dtype '{model_dtype}'") from exc
|
||||
self.model_dtype = model_dtype
|
||||
|
||||
attn_implementation = (
|
||||
"flash_attention_2" if (use_flash_attn and is_flash_attn_2_available()) else "eager"
|
||||
)
|
||||
attn_implementation = "flash_attention_2" if (use_flash_attn and _flash_attn_available()) else "eager"
|
||||
if use_flash_attn and attn_implementation == "eager":
|
||||
logger.warning(
|
||||
"Flash Attention 2 is unavailable on this runtime. Falling back to eager attention."
|
||||
)
|
||||
logger.warning("flash_attn is not installed. Falling back to eager attention.")
|
||||
|
||||
self.model = AutoModel.from_pretrained(
|
||||
model_name,
|
||||
@@ -365,3 +359,11 @@ class InternVL3Embedder(nn.Module):
|
||||
@property
|
||||
def device(self) -> torch.device:
|
||||
return next(self.model.parameters()).device
|
||||
|
||||
|
||||
def _flash_attn_available() -> bool:
|
||||
try:
|
||||
import flash_attn # noqa: F401
|
||||
except ModuleNotFoundError:
|
||||
return False
|
||||
return True
|
||||
|
||||
+318
-66
@@ -17,7 +17,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import inspect
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, TypedDict, Unpack
|
||||
|
||||
@@ -45,10 +44,26 @@ from lerobot.utils.constants import (
|
||||
)
|
||||
from lerobot.utils.feature_utils import dataset_to_policy_features
|
||||
|
||||
from .act.configuration_act import ACTConfig
|
||||
from .diffusion.configuration_diffusion import DiffusionConfig
|
||||
from .eo1.configuration_eo1 import EO1Config
|
||||
from .evo1.configuration_evo1 import Evo1Config
|
||||
from .fastwam.configuration_fastwam import FastWAMConfig
|
||||
from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig
|
||||
from .groot.configuration_groot import GrootConfig
|
||||
from .lingbot_va.configuration_lingbot_va import LingBotVAConfig
|
||||
from .molmoact2.configuration_molmoact2 import MolmoAct2Config
|
||||
from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
|
||||
from .pi0.configuration_pi0 import PI0Config
|
||||
from .pi05.configuration_pi05 import PI05Config
|
||||
from .pretrained import PreTrainedPolicy
|
||||
from .smolvla.configuration_smolvla import SmolVLAConfig
|
||||
from .tdmpc.configuration_tdmpc import TDMPCConfig
|
||||
from .utils import validate_visual_features_consistency
|
||||
from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig
|
||||
from .vqbet.configuration_vqbet import VQBeTConfig
|
||||
from .wall_x.configuration_wall_x import WallXConfig
|
||||
from .xvla.configuration_xvla import XVLAConfig
|
||||
|
||||
|
||||
def _reconnect_relative_absolute_steps(
|
||||
@@ -73,23 +88,100 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
|
||||
"""
|
||||
Retrieves a policy class by its registered name.
|
||||
|
||||
Resolution is convention-based: the draccus-registered config class of ``name`` is
|
||||
looked up, its ``configuration_*`` module path is rewritten to ``modeling_*``, and
|
||||
the ``<X>Policy`` class is imported from there. The modeling module is only imported
|
||||
at call time, keeping heavy optional dependencies lazy. This works for both built-in
|
||||
policies and third-party lerobot plugins (anything registered via
|
||||
``@PreTrainedConfig.register_subclass``).
|
||||
This function uses dynamic imports to avoid loading all policy classes into memory
|
||||
at once, improving startup time and reducing dependencies.
|
||||
|
||||
Args:
|
||||
name: The registered name of the policy (e.g. "act", "diffusion", "pi0").
|
||||
name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
|
||||
"multi_task_dit", "vqbet", "pi0", "pi05", "gaussian_actor", "smolvla", "wall_x",
|
||||
"molmoact2", "eo1", "evo1".
|
||||
Returns:
|
||||
The policy class corresponding to the given name.
|
||||
|
||||
Raises:
|
||||
ValueError: If the policy name is not registered.
|
||||
ImportError: If the policy's optional dependencies are not installed.
|
||||
NotImplementedError: If the policy name is not recognized.
|
||||
"""
|
||||
return _get_policy_cls_from_policy_name(name=name)
|
||||
if name == "tdmpc":
|
||||
from .tdmpc.modeling_tdmpc import TDMPCPolicy
|
||||
|
||||
return TDMPCPolicy
|
||||
elif name == "diffusion":
|
||||
from .diffusion.modeling_diffusion import DiffusionPolicy
|
||||
|
||||
return DiffusionPolicy
|
||||
elif name == "act":
|
||||
from .act.modeling_act import ACTPolicy
|
||||
|
||||
return ACTPolicy
|
||||
elif name == "multi_task_dit":
|
||||
from .multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy
|
||||
|
||||
return MultiTaskDiTPolicy
|
||||
elif name == "vqbet":
|
||||
from .vqbet.modeling_vqbet import VQBeTPolicy
|
||||
|
||||
return VQBeTPolicy
|
||||
elif name == "pi0":
|
||||
from .pi0.modeling_pi0 import PI0Policy
|
||||
|
||||
return PI0Policy
|
||||
elif name == "pi0_fast":
|
||||
from .pi0_fast.modeling_pi0_fast import PI0FastPolicy
|
||||
|
||||
return PI0FastPolicy
|
||||
elif name == "pi05":
|
||||
from .pi05.modeling_pi05 import PI05Policy
|
||||
|
||||
return PI05Policy
|
||||
elif name == "gaussian_actor":
|
||||
from .gaussian_actor.modeling_gaussian_actor import GaussianActorPolicy
|
||||
|
||||
return GaussianActorPolicy
|
||||
elif name == "smolvla":
|
||||
from .smolvla.modeling_smolvla import SmolVLAPolicy
|
||||
|
||||
return SmolVLAPolicy
|
||||
elif name == "groot":
|
||||
from .groot.modeling_groot import GrootPolicy
|
||||
|
||||
return GrootPolicy
|
||||
elif name == "xvla":
|
||||
from .xvla.modeling_xvla import XVLAPolicy
|
||||
|
||||
return XVLAPolicy
|
||||
elif name == "wall_x":
|
||||
from .wall_x.modeling_wall_x import WallXPolicy
|
||||
|
||||
return WallXPolicy
|
||||
elif name == "eo1":
|
||||
from .eo1.modeling_eo1 import EO1Policy
|
||||
|
||||
return EO1Policy
|
||||
elif name == "molmoact2":
|
||||
from .molmoact2.modeling_molmoact2 import MolmoAct2Policy
|
||||
|
||||
return MolmoAct2Policy
|
||||
elif name == "vla_jepa":
|
||||
from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy
|
||||
|
||||
return VLAJEPAPolicy
|
||||
elif name == "lingbot_va":
|
||||
from .lingbot_va.modeling_lingbot_va import LingBotVAPolicy
|
||||
|
||||
return LingBotVAPolicy
|
||||
elif name == "fastwam":
|
||||
from .fastwam.modeling_fastwam import FastWAMPolicy
|
||||
|
||||
return FastWAMPolicy
|
||||
elif name == "evo1":
|
||||
from .evo1.modeling_evo1 import Evo1Policy
|
||||
|
||||
return Evo1Policy
|
||||
else:
|
||||
try:
|
||||
return _get_policy_cls_from_policy_name(name=name)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Policy type '{name}' is not available.") from e
|
||||
|
||||
|
||||
def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
||||
@@ -100,8 +192,9 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
||||
mapping a string identifier to the corresponding config class.
|
||||
|
||||
Args:
|
||||
policy_type: The registered type of the policy (any name registered via
|
||||
``@PreTrainedConfig.register_subclass``, e.g. "act", "diffusion", "pi0").
|
||||
policy_type: The type of the policy. Supported types include "tdmpc",
|
||||
"multi_task_dit", "diffusion", "act", "vqbet", "pi0", "pi05", "gaussian_actor",
|
||||
"smolvla", "wall_x", "molmoact2", "eo1", "evo1".
|
||||
**kwargs: Keyword arguments to be passed to the configuration class constructor.
|
||||
|
||||
Returns:
|
||||
@@ -110,11 +203,48 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
||||
Raises:
|
||||
ValueError: If the `policy_type` is not recognized.
|
||||
"""
|
||||
try:
|
||||
config_cls = PreTrainedConfig.get_choice_class(policy_type)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Policy type '{policy_type}' is not available.") from e
|
||||
return config_cls(**kwargs)
|
||||
if policy_type == "tdmpc":
|
||||
return TDMPCConfig(**kwargs)
|
||||
elif policy_type == "diffusion":
|
||||
return DiffusionConfig(**kwargs)
|
||||
elif policy_type == "act":
|
||||
return ACTConfig(**kwargs)
|
||||
elif policy_type == "multi_task_dit":
|
||||
return MultiTaskDiTConfig(**kwargs)
|
||||
elif policy_type == "vqbet":
|
||||
return VQBeTConfig(**kwargs)
|
||||
elif policy_type == "pi0":
|
||||
return PI0Config(**kwargs)
|
||||
elif policy_type == "pi05":
|
||||
return PI05Config(**kwargs)
|
||||
elif policy_type == "gaussian_actor":
|
||||
return GaussianActorConfig(**kwargs)
|
||||
elif policy_type == "smolvla":
|
||||
return SmolVLAConfig(**kwargs)
|
||||
elif policy_type == "groot":
|
||||
return GrootConfig(**kwargs)
|
||||
elif policy_type == "xvla":
|
||||
return XVLAConfig(**kwargs)
|
||||
elif policy_type == "wall_x":
|
||||
return WallXConfig(**kwargs)
|
||||
elif policy_type == "eo1":
|
||||
return EO1Config(**kwargs)
|
||||
elif policy_type == "molmoact2":
|
||||
return MolmoAct2Config(**kwargs)
|
||||
elif policy_type == "vla_jepa":
|
||||
return VLAJEPAConfig(**kwargs)
|
||||
elif policy_type == "lingbot_va":
|
||||
return LingBotVAConfig(**kwargs)
|
||||
elif policy_type == "fastwam":
|
||||
return FastWAMConfig(**kwargs)
|
||||
elif policy_type == "evo1":
|
||||
return Evo1Config(**kwargs)
|
||||
else:
|
||||
try:
|
||||
config_cls = PreTrainedConfig.get_choice_class(policy_type)
|
||||
return config_cls(**kwargs)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Policy type '{policy_type}' is not available.") from e
|
||||
|
||||
|
||||
class ProcessorConfigKwargs(TypedDict, total=False):
|
||||
@@ -168,7 +298,8 @@ def make_pre_post_processors(
|
||||
A tuple containing the input (pre-processor) and output (post-processor) pipelines.
|
||||
|
||||
Raises:
|
||||
ValueError: If no processor factory exists for the given policy configuration type.
|
||||
NotImplementedError: If a processor factory is not implemented for the given
|
||||
policy configuration type.
|
||||
"""
|
||||
if pretrained_path:
|
||||
if isinstance(policy_cfg, GrootConfig):
|
||||
@@ -220,13 +351,166 @@ def make_pre_post_processors(
|
||||
)
|
||||
return preprocessor, postprocessor
|
||||
|
||||
# Create new processors from the policy config, resolving the per-policy factory
|
||||
# function by naming convention (lazy import keeps optional dependencies optional).
|
||||
return _make_processors_from_policy_config(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
dataset_meta=kwargs.get("dataset_meta"),
|
||||
)
|
||||
# Create a new processor based on policy type
|
||||
if isinstance(policy_cfg, TDMPCConfig):
|
||||
from .tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
|
||||
|
||||
processors = make_tdmpc_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, DiffusionConfig):
|
||||
from .diffusion.processor_diffusion import make_diffusion_pre_post_processors
|
||||
|
||||
processors = make_diffusion_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, ACTConfig):
|
||||
from .act.processor_act import make_act_pre_post_processors
|
||||
|
||||
processors = make_act_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, MultiTaskDiTConfig):
|
||||
from .multi_task_dit.processor_multi_task_dit import (
|
||||
make_multi_task_dit_pre_post_processors,
|
||||
)
|
||||
|
||||
processors = make_multi_task_dit_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, VQBeTConfig):
|
||||
from .vqbet.processor_vqbet import make_vqbet_pre_post_processors
|
||||
|
||||
processors = make_vqbet_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, PI0Config):
|
||||
from .pi0.processor_pi0 import make_pi0_pre_post_processors
|
||||
|
||||
processors = make_pi0_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, PI05Config):
|
||||
from .pi05.processor_pi05 import make_pi05_pre_post_processors
|
||||
|
||||
processors = make_pi05_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, GaussianActorConfig):
|
||||
from .gaussian_actor.processor_gaussian_actor import make_gaussian_actor_pre_post_processors
|
||||
|
||||
processors = make_gaussian_actor_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, SmolVLAConfig):
|
||||
from .smolvla.processor_smolvla import make_smolvla_pre_post_processors
|
||||
|
||||
processors = make_smolvla_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, GrootConfig):
|
||||
from .groot.processor_groot import make_groot_pre_post_processors
|
||||
|
||||
processors = make_groot_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
dataset_meta=kwargs.get("dataset_meta"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, XVLAConfig):
|
||||
from .xvla.processor_xvla import (
|
||||
make_xvla_pre_post_processors,
|
||||
)
|
||||
|
||||
processors = make_xvla_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, WallXConfig):
|
||||
from .wall_x.processor_wall_x import make_wall_x_pre_post_processors
|
||||
|
||||
processors = make_wall_x_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, EO1Config):
|
||||
from .eo1.processor_eo1 import make_eo1_pre_post_processors
|
||||
|
||||
processors = make_eo1_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
elif isinstance(policy_cfg, Evo1Config):
|
||||
from .evo1.processor_evo1 import make_evo1_pre_post_processors
|
||||
|
||||
processors = make_evo1_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, MolmoAct2Config):
|
||||
from .molmoact2.processor_molmoact2 import make_molmoact2_pre_post_processors
|
||||
|
||||
processors = make_molmoact2_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
dataset_meta=kwargs.get("dataset_meta"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, VLAJEPAConfig):
|
||||
from .vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors
|
||||
|
||||
processors = make_vla_jepa_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, LingBotVAConfig):
|
||||
from .lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors
|
||||
|
||||
processors = make_lingbot_va_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, FastWAMConfig):
|
||||
from .fastwam.processor_fastwam import make_fastwam_pre_post_processors
|
||||
|
||||
processors = make_fastwam_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
else:
|
||||
try:
|
||||
processors = _make_processors_from_policy_config(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Processor for policy type '{policy_cfg.type}' is not implemented.") from e
|
||||
|
||||
return processors
|
||||
|
||||
|
||||
def make_policy(
|
||||
@@ -370,12 +654,10 @@ def make_policy(
|
||||
return policy
|
||||
|
||||
|
||||
def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedPolicy]:
|
||||
def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedConfig]:
|
||||
"""Get policy class from its registered name using dynamic imports.
|
||||
|
||||
Works for built-in policies and 3rd party lerobot plugins alike: the config class
|
||||
registered under ``name`` is resolved via the draccus ChoiceRegistry, and the policy
|
||||
class is imported from the sibling ``modeling_*`` module by naming convention.
|
||||
This is used as a helper function to import policies from 3rd party lerobot plugins.
|
||||
|
||||
Args:
|
||||
name: The name of the policy.
|
||||
@@ -401,39 +683,22 @@ def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedPolicy]:
|
||||
"configuration_", "modeling_"
|
||||
) # e.g., configuration_diffusion -> modeling_diffusion
|
||||
|
||||
try:
|
||||
module = importlib.import_module(module_path)
|
||||
except ModuleNotFoundError as e:
|
||||
if e.name == module_path:
|
||||
# The modeling_* module itself does not exist for this policy type. A missing
|
||||
# optional dependency inside an existing module propagates unchanged instead,
|
||||
# so its actionable install hint stays visible.
|
||||
raise ValueError(f"Policy class for '{name}' is not implemented.") from e
|
||||
raise
|
||||
policy_cls = getattr(module, cls_name, None)
|
||||
if policy_cls is None:
|
||||
raise ValueError(
|
||||
f"Policy class '{cls_name}' not found in '{module_path}'. "
|
||||
f"Policies must expose '<Name>Policy' in the sibling 'modeling_*' module by naming convention."
|
||||
)
|
||||
module = importlib.import_module(module_path)
|
||||
policy_cls = getattr(module, cls_name)
|
||||
return policy_cls
|
||||
|
||||
|
||||
def _make_processors_from_policy_config(
|
||||
config: PreTrainedConfig,
|
||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||
dataset_meta: Any | None = None,
|
||||
) -> tuple[Any, Any]:
|
||||
"""Create pre- and post-processors from a policy configuration using dynamic imports.
|
||||
|
||||
Resolves ``make_{type}_pre_post_processors`` from the policy's ``processor_*`` module
|
||||
by naming convention. Works for built-in policies and 3rd party lerobot plugins.
|
||||
This is used as a helper function to import processor factories from 3rd party lerobot plugins.
|
||||
|
||||
Args:
|
||||
config: The policy configuration object.
|
||||
dataset_stats: Dataset statistics for normalization.
|
||||
dataset_meta: Dataset metadata, forwarded only to factories that declare a
|
||||
``dataset_meta`` parameter (e.g. groot, molmoact2).
|
||||
Returns:
|
||||
A tuple containing the input (pre-processor) and output (post-processor) pipelines.
|
||||
"""
|
||||
@@ -446,19 +711,6 @@ def _make_processors_from_policy_config(
|
||||
logging.debug(
|
||||
f"Instantiating pre/post processors using function '{function_name}' from module '{module_path}'"
|
||||
)
|
||||
try:
|
||||
module = importlib.import_module(module_path)
|
||||
except ModuleNotFoundError as e:
|
||||
if e.name == module_path:
|
||||
# The processor_* module itself does not exist for this policy type. A missing
|
||||
# optional dependency inside an existing module propagates unchanged instead,
|
||||
# so its actionable install hint stays visible.
|
||||
raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.") from e
|
||||
raise
|
||||
function = getattr(module, function_name, None)
|
||||
if function is None:
|
||||
raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.")
|
||||
call_kwargs: dict[str, Any] = {"dataset_stats": dataset_stats}
|
||||
if "dataset_meta" in inspect.signature(function).parameters:
|
||||
call_kwargs["dataset_meta"] = dataset_meta
|
||||
return function(config, **call_kwargs)
|
||||
module = importlib.import_module(module_path)
|
||||
function = getattr(module, function_name)
|
||||
return function(config, dataset_stats=dataset_stats)
|
||||
|
||||
@@ -22,11 +22,20 @@ import torch
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
ActionProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStepRegistry,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.utils.constants import (
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
|
||||
from .configuration_fastwam import FastWAMConfig
|
||||
@@ -96,20 +105,38 @@ def make_fastwam_pre_post_processors(
|
||||
# anyway) and unsafe across fine-tuning: its `resize_size` would be inherited from the base
|
||||
# checkpoint's camera geometry, not this dataset's, making the concatenation N_cameras x too wide.
|
||||
|
||||
steps = make_default_policy_processor_steps(config, normalization_stats, normalizer_device=config.device)
|
||||
|
||||
input_steps = [
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=normalization_stats,
|
||||
device=config.device,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
steps.unnormalize,
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=normalization_stats,
|
||||
),
|
||||
]
|
||||
if config.toggle_action_dimensions:
|
||||
output_steps.append(
|
||||
FastWAMActionToggleProcessorStep(toggle_dimensions=config.toggle_action_dimensions)
|
||||
)
|
||||
output_steps.append(steps.to_cpu)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
output_steps.append(DeviceProcessorStep(device="cpu"))
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -20,10 +20,17 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
make_default_pre_post_processors,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_gaussian_actor import GaussianActorConfig
|
||||
|
||||
@@ -55,4 +62,33 @@ def make_gaussian_actor_pre_post_processors(
|
||||
Returns:
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
return make_default_pre_post_processors(config, dataset_stats)
|
||||
|
||||
# Add remaining processors
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -25,12 +25,19 @@ import torch
|
||||
|
||||
from lerobot.configs.types import FeatureType, NormalizationMode
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
|
||||
from lerobot.utils.constants import (
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
|
||||
from .configuration_lingbot_va import LingBotVAConfig
|
||||
@@ -45,13 +52,15 @@ def make_lingbot_va_pre_post_processors(
|
||||
]:
|
||||
"""Build the pre/post processor pipelines for LingBot-VA."""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps: list[ProcessorStep] = [
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
steps.normalize,
|
||||
steps.to_device,
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
]
|
||||
|
||||
# Unnormalize actions from [-1, 1] to physical units (QUANTILES) using q01/q99 restored from the checkpoint.
|
||||
@@ -61,7 +70,18 @@ def make_lingbot_va_pre_post_processors(
|
||||
norm_map={FeatureType.ACTION: NormalizationMode.QUANTILES},
|
||||
stats=dataset_stats,
|
||||
),
|
||||
steps.to_cpu,
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -19,12 +19,18 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_multi_task_dit import MultiTaskDiTConfig
|
||||
|
||||
@@ -60,11 +66,9 @@ def make_multi_task_dit_pre_post_processors(
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats, normalizer_device=config.device)
|
||||
|
||||
input_steps = [
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name=config.text_encoder_name,
|
||||
padding=config.tokenizer_padding,
|
||||
@@ -72,12 +76,32 @@ def make_multi_task_dit_pre_post_processors(
|
||||
max_length=config.tokenizer_max_length,
|
||||
truncation=config.tokenizer_truncation,
|
||||
),
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
device=config.device,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
steps.unnormalize,
|
||||
steps.to_cpu,
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -21,16 +21,22 @@ import torch
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
ComplementaryDataProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RelativeActionsProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_pi0 import PI0Config
|
||||
|
||||
@@ -130,12 +136,10 @@ def make_pi0_pre_post_processors(
|
||||
action_names=getattr(config, "action_feature_names", None),
|
||||
)
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
# OpenPI order: raw → relative → normalize → model → unnormalize → absolute
|
||||
input_steps: list[ProcessorStep] = [
|
||||
steps.rename_observations, # To mimic the same processor as pretrained one
|
||||
steps.add_batch_dim,
|
||||
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
|
||||
AddBatchDimensionProcessorStep(),
|
||||
Pi0NewLineProcessor(), # Add newlines before tokenization for PaliGemma
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name="google/paligemma-3b-pt-224",
|
||||
@@ -143,15 +147,32 @@ def make_pi0_pre_post_processors(
|
||||
padding_side="right",
|
||||
padding="max_length",
|
||||
),
|
||||
steps.to_device,
|
||||
DeviceProcessorStep(device=config.device),
|
||||
relative_step,
|
||||
steps.normalize,
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
steps.unnormalize,
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
|
||||
steps.to_cpu,
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -24,17 +24,26 @@ import torch
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RelativeActionsProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
from lerobot.utils.constants import (
|
||||
OBS_STATE,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
|
||||
from .configuration_pi05 import PI05Config
|
||||
|
||||
@@ -126,16 +135,18 @@ def make_pi05_pre_post_processors(
|
||||
action_names=getattr(config, "action_feature_names", None),
|
||||
)
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
# OpenPI order: raw → relative → normalize → model → unnormalize → absolute
|
||||
input_steps: list[ProcessorStep] = [
|
||||
steps.rename_observations, # To mimic the same processor as pretrained one
|
||||
steps.add_batch_dim,
|
||||
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
|
||||
AddBatchDimensionProcessorStep(),
|
||||
relative_step,
|
||||
# NOTE: NormalizerProcessorStep MUST come before Pi05PrepareStateTokenizerProcessorStep
|
||||
# because the tokenizer step expects normalized state in [-1, 1] range for discretization
|
||||
steps.normalize,
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
Pi05PrepareStateTokenizerProcessorStep(max_state_dim=config.max_state_dim),
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name="google/paligemma-3b-pt-224",
|
||||
@@ -143,13 +154,26 @@ def make_pi05_pre_post_processors(
|
||||
padding_side="right",
|
||||
padding="max_length",
|
||||
),
|
||||
steps.to_device,
|
||||
DeviceProcessorStep(device=config.device),
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
steps.unnormalize,
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
|
||||
steps.to_cpu,
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -25,17 +25,26 @@ from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
ActionTokenizerProcessorStep,
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RelativeActionsProcessorStep,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
from lerobot.utils.constants import (
|
||||
OBS_STATE,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
|
||||
from .configuration_pi0_fast import PI0FastConfig
|
||||
|
||||
@@ -126,8 +135,6 @@ def make_pi0_fast_pre_post_processors(
|
||||
action_names=getattr(config, "action_feature_names", None),
|
||||
)
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
# Pi0Fast order: relative → normalize → tokenize → model → unnormalize → absolute
|
||||
# This matches pi0/pi0.5: RelativeActionsProcessorStep runs first on raw absolute actions,
|
||||
# caching the raw state. NormalizerProcessorStep then normalizes the raw relative actions,
|
||||
@@ -137,10 +144,14 @@ def make_pi0_fast_pre_post_processors(
|
||||
# before Pi0FastPrepareStateAndLanguageTokenizerProcessorStep, so the state tokenizer
|
||||
# continues to receive normalized state in [-1, 1] as expected.
|
||||
input_steps: list[ProcessorStep] = [
|
||||
steps.rename_observations, # To mimic the same processor as pretrained one
|
||||
steps.add_batch_dim,
|
||||
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
|
||||
AddBatchDimensionProcessorStep(),
|
||||
relative_step,
|
||||
steps.normalize,
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
Pi0FastPrepareStateAndLanguageTokenizerProcessorStep(max_state_dim=config.max_state_dim),
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name=config.text_tokenizer_name,
|
||||
@@ -154,13 +165,26 @@ def make_pi0_fast_pre_post_processors(
|
||||
fast_skip_tokens=config.fast_skip_tokens,
|
||||
paligemma_tokenizer_name=config.text_tokenizer_name,
|
||||
),
|
||||
steps.to_device,
|
||||
DeviceProcessorStep(device=config.device),
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
steps.unnormalize,
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
|
||||
steps.to_cpu,
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -23,6 +23,8 @@ from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack
|
||||
|
||||
import packaging
|
||||
import safetensors
|
||||
from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, save_torch_state_dict
|
||||
from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
|
||||
from huggingface_hub.errors import HfHubHTTPError
|
||||
@@ -32,7 +34,6 @@ from torch import Tensor, nn
|
||||
from lerobot.__version__ import __version__
|
||||
from lerobot.configs import PreTrainedConfig
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
from lerobot.utils.device_utils import resolve_safetensors_device
|
||||
from lerobot.utils.hub import HubMixin
|
||||
|
||||
from .utils import log_model_loading_keys
|
||||
@@ -220,10 +221,26 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
|
||||
@classmethod
|
||||
def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
|
||||
missing_keys, unexpected_keys = load_model_as_safetensor(
|
||||
model, model_file, strict=strict, device=resolve_safetensors_device(map_location)
|
||||
)
|
||||
# Create base kwargs
|
||||
kwargs = {"strict": strict}
|
||||
|
||||
# Add device parameter for newer versions that support it
|
||||
if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
|
||||
kwargs["device"] = map_location
|
||||
|
||||
# Load the model with appropriate kwargs
|
||||
missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
|
||||
log_model_loading_keys(missing_keys, unexpected_keys)
|
||||
|
||||
# For older versions, manually move to device if needed
|
||||
if "device" not in kwargs and map_location != "cpu":
|
||||
logging.warning(
|
||||
"Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
|
||||
" This means that the model is loaded on 'cpu' first and then copied to the device."
|
||||
" This leads to a slower loading time."
|
||||
" Please update safetensors to version 0.4.3 or above for improved performance."
|
||||
)
|
||||
model.to(map_location)
|
||||
return model
|
||||
|
||||
@abc.abstractmethod
|
||||
|
||||
@@ -19,13 +19,19 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NewLineTaskProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_smolvla import SmolVLAConfig
|
||||
|
||||
@@ -60,11 +66,9 @@ def make_smolvla_pre_post_processors(
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps = [
|
||||
steps.rename_observations, # To mimic the same processor as pretrained one
|
||||
steps.add_batch_dim,
|
||||
RenameObservationsProcessorStep(rename_map={}), # To mimic the same processor as pretrained one
|
||||
AddBatchDimensionProcessorStep(),
|
||||
NewLineTaskProcessorStep(),
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name=config.vlm_model_name,
|
||||
@@ -72,11 +76,28 @@ def make_smolvla_pre_post_processors(
|
||||
padding_side="right",
|
||||
max_length=config.tokenizer_max_length,
|
||||
),
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
steps.unnormalize,
|
||||
steps.to_cpu,
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -19,10 +19,17 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
make_default_pre_post_processors,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_tdmpc import TDMPCConfig
|
||||
|
||||
@@ -54,4 +61,32 @@ def make_tdmpc_pre_post_processors(
|
||||
Returns:
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
return make_default_pre_post_processors(config, dataset_stats)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -20,16 +20,20 @@ import torch
|
||||
|
||||
from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
EnvTransition,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
TransitionKey,
|
||||
UnnormalizerProcessorStep,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="vla_jepa_clip_actions")
|
||||
@@ -108,12 +112,15 @@ def make_vla_jepa_pre_post_processors(
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
]:
|
||||
features = {**config.input_features, **config.output_features}
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
input_steps = [
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features=features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps: list[ProcessorStep] = []
|
||||
if config.clip_normalized_actions:
|
||||
@@ -122,8 +129,6 @@ def make_vla_jepa_pre_post_processors(
|
||||
output_steps.append(
|
||||
PreSnapGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)
|
||||
)
|
||||
# NOTE: unlike the default policy unnormalizer (output features only), VLA-JEPA
|
||||
# unnormalizes over BOTH input and output features.
|
||||
output_steps.append(
|
||||
UnnormalizerProcessorStep(
|
||||
features=features,
|
||||
@@ -135,5 +140,16 @@ def make_vla_jepa_pre_post_processors(
|
||||
output_steps.append(
|
||||
BinarizeGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)
|
||||
)
|
||||
output_steps.append(steps.to_cpu)
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
output_steps.append(DeviceProcessorStep(device="cpu"))
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -20,10 +20,17 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
make_default_pre_post_processors,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_vqbet import VQBeTConfig
|
||||
|
||||
@@ -55,4 +62,32 @@ def make_vqbet_pre_post_processors(
|
||||
Returns:
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
"""
|
||||
return make_default_pre_post_processors(config, dataset_stats)
|
||||
|
||||
input_steps = [
|
||||
RenameObservationsProcessorStep(rename_map={}), # Let the possibility to the user to rename the keys
|
||||
AddBatchDimensionProcessorStep(),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -20,13 +20,19 @@ import torch
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
ComplementaryDataProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStepRegistry,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .configuration_wall_x import WallXConfig
|
||||
|
||||
@@ -59,22 +65,37 @@ def make_wall_x_pre_post_processors(
|
||||
A tuple containing the configured pre-processor and post-processor pipelines
|
||||
"""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
input_steps = [
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
WallXTaskProcessor(), # Process task description
|
||||
steps.normalize,
|
||||
steps.to_device,
|
||||
NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
]
|
||||
|
||||
output_steps = [
|
||||
steps.unnormalize,
|
||||
steps.to_cpu,
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="wall_x_task_processor")
|
||||
|
||||
@@ -22,14 +22,19 @@ import torch
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
ObservationProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
TokenizerProcessorStep,
|
||||
make_default_policy_processor_steps,
|
||||
make_policy_processor_pipelines,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
@@ -37,6 +42,8 @@ from lerobot.utils.constants import (
|
||||
OBS_IMAGES,
|
||||
OBS_PREFIX,
|
||||
OBS_STATE,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
|
||||
from .configuration_xvla import XVLAConfig
|
||||
@@ -54,11 +61,10 @@ def make_xvla_pre_post_processors(
|
||||
Build the LeRobot processor pipelines for XVLA.
|
||||
"""
|
||||
|
||||
steps = make_default_policy_processor_steps(config, dataset_stats)
|
||||
|
||||
features = {**config.input_features, **config.output_features}
|
||||
input_steps = [
|
||||
steps.rename_observations,
|
||||
steps.add_batch_dim,
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name=config.tokenizer_name,
|
||||
max_length=config.tokenizer_max_length,
|
||||
@@ -68,15 +74,32 @@ def make_xvla_pre_post_processors(
|
||||
XVLAImageToFloatProcessorStep(),
|
||||
XVLAImageNetNormalizeProcessorStep(),
|
||||
XVLAAddDomainIdProcessorStep(),
|
||||
steps.to_device,
|
||||
steps.normalize,
|
||||
DeviceProcessorStep(device=config.device),
|
||||
NormalizerProcessorStep(
|
||||
features=features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
]
|
||||
output_steps = [
|
||||
steps.unnormalize,
|
||||
steps.to_cpu,
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# Custom XVLA processor steps
|
||||
|
||||
@@ -42,14 +42,10 @@ from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorTo
|
||||
from .device_processor import DeviceProcessorStep
|
||||
from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep
|
||||
from .factory import (
|
||||
DefaultPolicyProcessorSteps,
|
||||
make_default_policy_processor_steps,
|
||||
make_default_pre_post_processors,
|
||||
make_default_processors,
|
||||
make_default_robot_action_processor,
|
||||
make_default_robot_observation_processor,
|
||||
make_default_teleop_action_processor,
|
||||
make_policy_processor_pipelines,
|
||||
)
|
||||
from .gym_action_processor import (
|
||||
Numpy2TorchActionProcessorStep,
|
||||
@@ -133,14 +129,10 @@ __all__ = [
|
||||
"ImageCropResizeProcessorStep",
|
||||
"InfoProcessorStep",
|
||||
"InterventionActionProcessorStep",
|
||||
"DefaultPolicyProcessorSteps",
|
||||
"make_default_policy_processor_steps",
|
||||
"make_default_pre_post_processors",
|
||||
"make_default_processors",
|
||||
"make_default_teleop_action_processor",
|
||||
"make_default_robot_action_processor",
|
||||
"make_default_robot_observation_processor",
|
||||
"make_policy_processor_pipelines",
|
||||
"AbsoluteActionsProcessorStep",
|
||||
"RelativeActionsProcessorStep",
|
||||
"MapDeltaActionToRobotActionStep",
|
||||
|
||||
@@ -14,33 +14,15 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from lerobot.types import RobotAction, RobotObservation
|
||||
|
||||
import torch
|
||||
|
||||
from lerobot.configs.policies import PreTrainedConfig
|
||||
from lerobot.types import PolicyAction, RobotAction, RobotObservation
|
||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||
|
||||
from .batch_processor import AddBatchDimensionProcessorStep
|
||||
from .converters import (
|
||||
observation_to_transition,
|
||||
policy_action_to_transition,
|
||||
robot_action_observation_to_transition,
|
||||
transition_to_observation,
|
||||
transition_to_policy_action,
|
||||
transition_to_robot_action,
|
||||
)
|
||||
from .device_processor import DeviceProcessorStep
|
||||
from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
|
||||
from .pipeline import (
|
||||
IdentityProcessorStep,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
RobotProcessorPipeline,
|
||||
)
|
||||
from .rename_processor import RenameObservationsProcessorStep
|
||||
from .pipeline import IdentityProcessorStep, RobotProcessorPipeline
|
||||
|
||||
|
||||
def make_default_teleop_action_processor() -> RobotProcessorPipeline[
|
||||
@@ -79,97 +61,3 @@ def make_default_processors():
|
||||
robot_action_processor = make_default_robot_action_processor()
|
||||
robot_observation_processor = make_default_robot_observation_processor()
|
||||
return (teleop_action_processor, robot_action_processor, robot_observation_processor)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DefaultPolicyProcessorSteps:
|
||||
"""The canonical processor steps shared by most policies' pre/post pipelines.
|
||||
|
||||
Policies compose these in their own order (step ORDER is a Hub-serialized contract
|
||||
and intentionally stays explicit per policy) and interleave their custom steps.
|
||||
"""
|
||||
|
||||
rename_observations: RenameObservationsProcessorStep
|
||||
add_batch_dim: AddBatchDimensionProcessorStep
|
||||
to_device: DeviceProcessorStep
|
||||
normalize: NormalizerProcessorStep
|
||||
unnormalize: UnnormalizerProcessorStep
|
||||
to_cpu: DeviceProcessorStep
|
||||
|
||||
|
||||
def make_default_policy_processor_steps(
|
||||
config: PreTrainedConfig,
|
||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||
*,
|
||||
normalizer_device: torch.device | str | None = None,
|
||||
) -> DefaultPolicyProcessorSteps:
|
||||
"""Construct the canonical policy processor steps from a policy config.
|
||||
|
||||
Args:
|
||||
config: A `PreTrainedConfig` providing `device`, `input_features`,
|
||||
`output_features` and `normalization_mapping`.
|
||||
dataset_stats: Dataset statistics used for (un)normalization.
|
||||
normalizer_device: Device passed to `NormalizerProcessorStep` (some policies pin
|
||||
their normalization stats to the policy device; most leave it unset).
|
||||
"""
|
||||
return DefaultPolicyProcessorSteps(
|
||||
rename_observations=RenameObservationsProcessorStep(rename_map={}),
|
||||
add_batch_dim=AddBatchDimensionProcessorStep(),
|
||||
to_device=DeviceProcessorStep(device=config.device),
|
||||
normalize=NormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
device=normalizer_device,
|
||||
),
|
||||
unnormalize=UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
),
|
||||
to_cpu=DeviceProcessorStep(device="cpu"),
|
||||
)
|
||||
|
||||
|
||||
def make_policy_processor_pipelines(
|
||||
input_steps: list[ProcessorStep],
|
||||
output_steps: list[ProcessorStep],
|
||||
) -> tuple[
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
]:
|
||||
"""Wrap pre/post step lists into the canonical policy pipeline pair.
|
||||
|
||||
Uses the standard pipeline names (which determine the serialized JSON filenames on
|
||||
the Hub) and the standard policy-action converters on the postprocessor.
|
||||
"""
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def make_default_pre_post_processors(
|
||||
config: PreTrainedConfig,
|
||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||
*,
|
||||
normalizer_device: torch.device | str | None = None,
|
||||
) -> tuple[
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
]:
|
||||
"""The pure-scaffold policy pipeline pair: Rename -> Batch -> Device -> Normalize,
|
||||
and Unnormalize -> Device(cpu). Policies with custom steps or a different step order
|
||||
compose `make_default_policy_processor_steps` themselves instead.
|
||||
"""
|
||||
s = make_default_policy_processor_steps(config, dataset_stats, normalizer_device=normalizer_device)
|
||||
return make_policy_processor_pipelines(
|
||||
input_steps=[s.rename_observations, s.add_batch_dim, s.to_device, s.normalize],
|
||||
output_steps=[s.unnormalize, s.to_cpu],
|
||||
)
|
||||
|
||||
@@ -126,7 +126,7 @@ class RelativeActionsProcessorStep(ProcessorStep):
|
||||
observation = transition.get(TransitionKey.OBSERVATION, {})
|
||||
state = observation.get(OBS_STATE) if observation else None
|
||||
|
||||
# Always cache state for the paired AbsoluteActionsProcessorStep.
|
||||
# Always cache state for the paired AbsoluteActionsProcessorStep
|
||||
if state is not None:
|
||||
self._last_state = state
|
||||
|
||||
@@ -146,11 +146,6 @@ class RelativeActionsProcessorStep(ProcessorStep):
|
||||
"""Return the cached ``observation.state`` used as the reference point for relative/absolute action conversions."""
|
||||
return self._last_state
|
||||
|
||||
def set_cached_state(self, state: torch.Tensor | None) -> None:
|
||||
"""Override the cached anchor state, e.g. to re-pin a chunk's anchor after the
|
||||
per-tick pipeline overwrote it (see ``SyncInferenceEngine``)."""
|
||||
self._last_state = state
|
||||
|
||||
def get_config(self) -> dict[str, Any]:
|
||||
return {
|
||||
"enabled": self.enabled,
|
||||
|
||||
@@ -21,6 +21,8 @@ from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
from typing import TYPE_CHECKING, Any, TypeVar
|
||||
|
||||
import packaging
|
||||
import safetensors
|
||||
from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download
|
||||
from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
|
||||
from huggingface_hub.errors import HfHubHTTPError
|
||||
@@ -28,7 +30,6 @@ from safetensors.torch import load_model as load_model_as_safetensor, save_model
|
||||
from torch import Tensor, nn
|
||||
|
||||
from lerobot.configs.rewards import RewardModelConfig
|
||||
from lerobot.utils.device_utils import resolve_safetensors_device
|
||||
from lerobot.utils.hub import HubMixin
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -128,13 +129,29 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
|
||||
|
||||
@classmethod
|
||||
def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
|
||||
missing_keys, unexpected_keys = load_model_as_safetensor(
|
||||
model, model_file, strict=strict, device=resolve_safetensors_device(map_location)
|
||||
)
|
||||
# Create base kwargs
|
||||
kwargs = {"strict": strict}
|
||||
|
||||
# Add device parameter for newer versions that support it
|
||||
if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
|
||||
kwargs["device"] = map_location
|
||||
|
||||
# Load the model with appropriate kwargs
|
||||
missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
|
||||
if missing_keys:
|
||||
logging.warning(f"Missing key(s) when loading model: {missing_keys}")
|
||||
if unexpected_keys:
|
||||
logging.warning(f"Unexpected key(s) when loading model: {unexpected_keys}")
|
||||
|
||||
# For older versions, manually move to device if needed
|
||||
if "device" not in kwargs and map_location != "cpu":
|
||||
logging.warning(
|
||||
"Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
|
||||
" This means that the model is loaded on 'cpu' first and then copied to the device."
|
||||
" This leads to a slower loading time."
|
||||
" Please update safetensors to version 0.4.3 or above for improved performance."
|
||||
)
|
||||
model.to(map_location)
|
||||
return model
|
||||
|
||||
def get_optim_params(self):
|
||||
|
||||
@@ -14,7 +14,6 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
@@ -34,8 +33,6 @@ from lerobot.processor import (
|
||||
)
|
||||
from lerobot.utils.rotation import Rotation
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register("ee_reference_and_delta")
|
||||
@dataclass
|
||||
@@ -197,17 +194,11 @@ class EEBoundsAndSafety(RobotActionProcessorStep):
|
||||
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.
|
||||
"""
|
||||
|
||||
end_effector_bounds: dict
|
||||
max_ee_step_m: float = 0.05
|
||||
raise_on_jump: bool = True
|
||||
_last_pos: np.ndarray | None = field(default=None, init=False, repr=False)
|
||||
|
||||
def action(self, action: RobotAction) -> RobotAction:
|
||||
@@ -235,19 +226,8 @@ class EEBoundsAndSafety(RobotActionProcessorStep):
|
||||
dpos = pos - self._last_pos
|
||||
n = float(np.linalg.norm(dpos))
|
||||
if n > self.max_ee_step_m and n > 0:
|
||||
# Clamp the step to the per-frame limit (rate-limit). The clamped
|
||||
# value is computed either way; raise_on_jump only decides whether
|
||||
# an over-limit step aborts the loop or is rate-limited + warned.
|
||||
pos = self._last_pos + dpos * (self.max_ee_step_m / n)
|
||||
if self.raise_on_jump:
|
||||
raise ValueError(f"EE jump {n:.3f}m > {self.max_ee_step_m}m")
|
||||
logger.warning(
|
||||
"EE jump %.3fm > %.3fm; rate-limited to the per-frame step "
|
||||
"(likely a transient tracking glitch; if it recurs every frame "
|
||||
"the commanded target is systematically out of workspace).",
|
||||
n,
|
||||
self.max_ee_step_m,
|
||||
)
|
||||
raise ValueError(f"EE jump {n:.3f}m > {self.max_ee_step_m}m")
|
||||
|
||||
self._last_pos = pos
|
||||
|
||||
@@ -284,18 +264,12 @@ class InverseKinematicsEEToJoints(RobotActionProcessorStep):
|
||||
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).
|
||||
"""
|
||||
|
||||
kinematics: RobotKinematics
|
||||
motor_names: list[str]
|
||||
q_curr: np.ndarray | None = field(default=None, init=False, repr=False)
|
||||
initial_guess_current_joints: bool = True
|
||||
orientation_weight: float = 0.01
|
||||
|
||||
def action(self, action: RobotAction) -> RobotAction:
|
||||
x = action.pop("ee.x")
|
||||
@@ -334,9 +308,7 @@ class InverseKinematicsEEToJoints(RobotActionProcessorStep):
|
||||
t_des[:3, 3] = [x, y, z]
|
||||
|
||||
# Compute inverse kinematics
|
||||
q_target = self.kinematics.inverse_kinematics(
|
||||
self.q_curr, t_des, orientation_weight=self.orientation_weight
|
||||
)
|
||||
q_target = self.kinematics.inverse_kinematics(self.q_curr, t_des)
|
||||
self.q_curr = q_target
|
||||
|
||||
# TODO: This is sentitive to order of motor_names = q_target mapping
|
||||
|
||||
@@ -43,6 +43,7 @@ from lerobot.processor import (
|
||||
make_default_processors,
|
||||
rename_stats,
|
||||
)
|
||||
from lerobot.processor.relative_action_processor import RelativeActionsProcessorStep
|
||||
from lerobot.robots import make_robot_from_config
|
||||
from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
|
||||
from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
|
||||
@@ -51,6 +52,7 @@ from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
|
||||
from .inference import (
|
||||
InferenceEngine,
|
||||
RTCInferenceConfig,
|
||||
SyncInferenceConfig,
|
||||
create_inference_engine,
|
||||
)
|
||||
from .robot_wrapper import ThreadSafeRobot
|
||||
@@ -397,6 +399,15 @@ def build_rollout_context(
|
||||
},
|
||||
)
|
||||
|
||||
if isinstance(cfg.inference, SyncInferenceConfig) and any(
|
||||
isinstance(step, RelativeActionsProcessorStep) and step.enabled
|
||||
for step in getattr(preprocessor, "steps", ())
|
||||
):
|
||||
raise NotImplementedError(
|
||||
"SyncInferenceEngine does not support policies with relative actions for now."
|
||||
"Use --inference.type=rtc or remove relative action processor steps from the policy pipeline."
|
||||
)
|
||||
|
||||
# --- 7. Inference strategy (needs policy + pre/post + hardware) --
|
||||
logger.info(
|
||||
"Creating inference engine (type=%s)...",
|
||||
|
||||
@@ -24,21 +24,26 @@ import torch
|
||||
|
||||
from lerobot.policies.pretrained import PreTrainedPolicy
|
||||
from lerobot.policies.utils import make_robot_action, prepare_observation_for_inference
|
||||
from lerobot.processor import PolicyProcessorPipeline, RelativeActionsProcessorStep
|
||||
from lerobot.processor import PolicyProcessorPipeline
|
||||
|
||||
from .base import InferenceEngine
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Relative-action support: a predicted chunk of offsets is anchored to the robot
|
||||
# state at prediction time, but the sync engine reruns the pre/post pipeline every
|
||||
# tick, so ``RelativeActionsProcessorStep`` would re-anchor cached actions to the
|
||||
# current (moved) state and drift through the chunk. We pin the anchor per chunk:
|
||||
# a probe on the policy's public ``predict_action_chunk`` flags the ticks that
|
||||
# predict a fresh chunk; on the others the engine restores the anchor the relative
|
||||
# step overwrote. ``select_action`` stays on the hot path, so per-tick side effects
|
||||
# (e.g. LingBot-VA keyframe feedback) are preserved.
|
||||
# TODO(Steven): support relative-action policies. The per-tick flow refreshes
|
||||
# ``RelativeActionsProcessorStep._last_state`` every call, so cached chunk
|
||||
# actions popped on later ticks get reanchored to the *current* robot state and
|
||||
# absolute targets drift through the chunk. Relative-action policies are
|
||||
# rejected at context-build time today; RTC postprocesses the whole chunk and
|
||||
# is unaffected.
|
||||
#
|
||||
# Candidate fix: drive the policy via ``predict_action_chunk`` and serve a
|
||||
# local FIFO of postprocessed actions. Eliminates drift by construction and
|
||||
# saves per-tick pre/post work, but bypasses ``select_action`` — needs
|
||||
# fallbacks for SAC (raises), ACT temporal ensembling (ensembler lives in
|
||||
# ``select_action``), and Diffusion-family (obs-history queues populated as a
|
||||
# side effect of ``select_action``).
|
||||
|
||||
|
||||
class SyncInferenceEngine(InferenceEngine):
|
||||
@@ -68,31 +73,6 @@ class SyncInferenceEngine(InferenceEngine):
|
||||
self._task = task
|
||||
self._device = torch.device(device or "cpu")
|
||||
self._robot_type = robot_type
|
||||
|
||||
# Find an enabled RelativeActionsProcessorStep to pin its anchor per chunk
|
||||
# (see module comment), mirroring the RTC engine.
|
||||
self._relative_step = next(
|
||||
(
|
||||
s
|
||||
for s in getattr(preprocessor, "steps", ())
|
||||
if isinstance(s, RelativeActionsProcessorStep) and s.enabled
|
||||
),
|
||||
None,
|
||||
)
|
||||
# Set by the probe for the current tick / ever, respectively.
|
||||
self._chunk_predicted = False
|
||||
self._ever_predicted_chunk = False
|
||||
self._original_predict_action_chunk = None # set while the probe is installed
|
||||
if self._relative_step is not None:
|
||||
# ``action_names`` is optional on the step; fill it lazily from the
|
||||
# policy/dataset so the relative<->absolute mask is built correctly. This is
|
||||
# a deliberate engine->step side effect (the step is configured by its consumer).
|
||||
if self._relative_step.action_names is None:
|
||||
cfg_names = getattr(policy.config, "action_feature_names", None)
|
||||
self._relative_step.action_names = list(cfg_names) if cfg_names else list(ordered_action_keys)
|
||||
self._install_chunk_probe()
|
||||
logger.info("Relative actions enabled: chunk anchor pinned per predicted chunk")
|
||||
|
||||
logger.info(
|
||||
"SyncInferenceEngine initialized (device=%s, action_keys=%d)",
|
||||
self._device,
|
||||
@@ -105,11 +85,6 @@ class SyncInferenceEngine(InferenceEngine):
|
||||
|
||||
def stop(self) -> None:
|
||||
"""No background resources to stop."""
|
||||
# Undo the probe so the policy object isn't left permanently patched
|
||||
# (it may outlive this engine or be reused by another).
|
||||
if self._original_predict_action_chunk is not None:
|
||||
self._policy.predict_action_chunk = self._original_predict_action_chunk
|
||||
self._original_predict_action_chunk = None
|
||||
logger.info("SyncInferenceEngine stopped")
|
||||
|
||||
def reset(self) -> None:
|
||||
@@ -118,27 +93,6 @@ class SyncInferenceEngine(InferenceEngine):
|
||||
self._policy.reset()
|
||||
self._preprocessor.reset()
|
||||
self._postprocessor.reset()
|
||||
# New episode: the next tick predicts a fresh chunk and re-anchors.
|
||||
self._chunk_predicted = False
|
||||
self._ever_predicted_chunk = False
|
||||
|
||||
def _install_chunk_probe(self) -> None:
|
||||
"""Wrap the policy's public ``predict_action_chunk`` so we learn which ticks
|
||||
predict a fresh chunk (when the anchor must advance) without introspecting any
|
||||
private action queue. Chunking policies call it from ``select_action``.
|
||||
|
||||
Wraps whatever callable is currently bound (e.g. an already-``torch.compile``d
|
||||
one, since ``build_rollout_context`` compiles before building the engine); undone
|
||||
in ``stop()``."""
|
||||
self._original_predict_action_chunk = self._policy.predict_action_chunk
|
||||
inner = self._original_predict_action_chunk
|
||||
|
||||
def probe(*args, **kwargs):
|
||||
self._chunk_predicted = True
|
||||
self._ever_predicted_chunk = True
|
||||
return inner(*args, **kwargs)
|
||||
|
||||
self._policy.predict_action_chunk = probe
|
||||
|
||||
def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
|
||||
"""Run the full inference pipeline on ``obs_frame`` and return an action tensor."""
|
||||
@@ -153,25 +107,12 @@ class SyncInferenceEngine(InferenceEngine):
|
||||
if self._device.type == "cuda" and self._policy.config.use_amp
|
||||
else nullcontext()
|
||||
)
|
||||
# Snapshot the chunk anchor before the preprocessor overwrites it with this
|
||||
# tick's state; restore it below if this tick only served a cached action.
|
||||
# ``clone`` so the snapshot survives even if the cached tensor is ever mutated
|
||||
# in place (today it is only rebound, but the copy is cheap for a state vector).
|
||||
anchor_before = None
|
||||
if self._relative_step is not None:
|
||||
cached = self._relative_step.get_cached_state()
|
||||
anchor_before = cached.clone() if cached is not None else None
|
||||
self._chunk_predicted = False
|
||||
with torch.inference_mode(), autocast_ctx:
|
||||
observation = prepare_observation_for_inference(
|
||||
observation, self._device, self._task, self._robot_type
|
||||
)
|
||||
observation = self._preprocessor(observation)
|
||||
action = self._policy.select_action(observation)
|
||||
# Hold the anchor only for a chunking policy serving a cached action this
|
||||
# tick; policies that never chunk or that recomputed keep refreshing.
|
||||
if self._relative_step is not None and self._ever_predicted_chunk and not self._chunk_predicted:
|
||||
self._relative_step.set_cached_state(anchor_before)
|
||||
action = self._postprocessor(action)
|
||||
action_tensor = action.squeeze(0).cpu()
|
||||
|
||||
|
||||
@@ -28,12 +28,7 @@ For distributed runs, see ``examples/annotations/run_hf_job.py``.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from huggingface_hub import HfApi, snapshot_download
|
||||
from huggingface_hub.errors import RevisionNotFoundError
|
||||
|
||||
from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
|
||||
from lerobot.annotations.steerable_pipeline.executor import Executor
|
||||
@@ -47,12 +42,6 @@ from lerobot.annotations.steerable_pipeline.validator import StagingValidator
|
||||
from lerobot.annotations.steerable_pipeline.vlm_client import make_vlm_client
|
||||
from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter
|
||||
from lerobot.configs import parser
|
||||
from lerobot.utils.import_utils import _datasets_available, require_package
|
||||
|
||||
if TYPE_CHECKING or _datasets_available:
|
||||
from lerobot.datasets.dataset_metadata import CODEBASE_VERSION
|
||||
from lerobot.datasets.io_utils import load_info
|
||||
from lerobot.datasets.utils import create_lerobot_dataset_card
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -61,6 +50,8 @@ def _resolve_root(cfg: AnnotationPipelineConfig) -> Path:
|
||||
if cfg.root is not None:
|
||||
return Path(cfg.root)
|
||||
if cfg.repo_id is not None:
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
return Path(snapshot_download(repo_id=cfg.repo_id, repo_type="dataset"))
|
||||
raise ValueError("Either --root or --repo_id must be provided.")
|
||||
|
||||
@@ -134,7 +125,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
||||
|
||||
Pushes to ``cfg.new_repo_id`` when set, otherwise back to ``cfg.repo_id``.
|
||||
"""
|
||||
require_package("datasets", "dataset")
|
||||
from huggingface_hub import HfApi # noqa: PLC0415
|
||||
|
||||
repo_id = cfg.new_repo_id or cfg.repo_id
|
||||
commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
|
||||
@@ -152,26 +143,33 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
||||
repo_id=repo_id,
|
||||
repo_type="dataset",
|
||||
commit_message=commit_message,
|
||||
# README.md is excluded because when pushing to ``new_repo_id`` the
|
||||
# source card's links (e.g. the visualize badge) would keep pointing
|
||||
# at the source dataset; a fresh card is generated below instead.
|
||||
ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
|
||||
ignore_patterns=[".annotate_staging/**", "**/.DS_Store"],
|
||||
)
|
||||
print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True)
|
||||
|
||||
dataset_info = load_info(root)
|
||||
card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
|
||||
card.push_to_hub(repo_id=repo_id, repo_type="dataset")
|
||||
|
||||
# Tag the upload with the codebase version. ``LeRobotDatasetMetadata``
|
||||
# resolves the dataset revision via ``get_safe_version`` which scans
|
||||
# for tags like ``v3.0``; without a tag it raises
|
||||
# ``RevisionNotFoundError``. Read the version straight from the
|
||||
# dataset's own ``meta/info.json`` so we tag whatever the writer
|
||||
# actually wrote (no accidental drift if the codebase floor moves).
|
||||
version_tag = (
|
||||
dataset_info.codebase_version if dataset_info.codebase_version.startswith("v") else CODEBASE_VERSION
|
||||
)
|
||||
from lerobot.datasets.dataset_metadata import CODEBASE_VERSION # noqa: PLC0415
|
||||
|
||||
info_path = root / "meta" / "info.json"
|
||||
version_tag = CODEBASE_VERSION
|
||||
if info_path.exists():
|
||||
try:
|
||||
from lerobot.utils.io_utils import load_json # noqa: PLC0415
|
||||
|
||||
info = load_json(info_path)
|
||||
ds_version = info.get("codebase_version")
|
||||
if isinstance(ds_version, str) and ds_version.startswith("v"):
|
||||
version_tag = ds_version
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(
|
||||
f"[lerobot-annotate] could not read codebase_version from info.json ({exc}); falling back to {version_tag}",
|
||||
flush=True,
|
||||
)
|
||||
revision = getattr(commit_info, "oid", None)
|
||||
tag_kwargs = {
|
||||
"repo_id": repo_id,
|
||||
@@ -182,6 +180,10 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
||||
tag_kwargs["revision"] = revision
|
||||
|
||||
try:
|
||||
from contextlib import suppress # noqa: PLC0415
|
||||
|
||||
from huggingface_hub.errors import RevisionNotFoundError # noqa: PLC0415
|
||||
|
||||
with suppress(RevisionNotFoundError):
|
||||
api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
|
||||
api.create_tag(**tag_kwargs)
|
||||
|
||||
@@ -171,9 +171,6 @@ def update_policy(
|
||||
train_metrics.update_s = time.perf_counter() - start_time
|
||||
if torch.cuda.is_available():
|
||||
train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
|
||||
# Aggregate the policy's scalar outputs for logging and rank-reduction across the log window.
|
||||
if output_dict:
|
||||
train_metrics.update_metrics(output_dict)
|
||||
return train_metrics, output_dict
|
||||
|
||||
|
||||
@@ -575,7 +572,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
batch = preprocessor(batch)
|
||||
train_tracker.dataloading_s = time.perf_counter() - start_time
|
||||
|
||||
train_tracker, _ = update_policy(
|
||||
train_tracker, output_dict = update_policy(
|
||||
train_tracker,
|
||||
policy,
|
||||
batch,
|
||||
@@ -608,10 +605,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
train_tracker.samples_per_s = effective_batch_size / step_time
|
||||
logging.info(train_tracker)
|
||||
if wandb_logger:
|
||||
# Policy sub-losses (latent_loss, action_loss, ...) are aggregated into the
|
||||
# tracker by update_policy, so to_dict() already carries their windowed,
|
||||
# rank-reduced averages — no per-step output_dict passthrough needed.
|
||||
wandb_log_dict = train_tracker.to_dict()
|
||||
if output_dict:
|
||||
wandb_log_dict.update(output_dict)
|
||||
# Log sample weighting statistics if enabled
|
||||
if sample_weighter is not None:
|
||||
weighter_stats = sample_weighter.get_stats()
|
||||
|
||||
@@ -59,20 +59,6 @@ def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
|
||||
return device
|
||||
|
||||
|
||||
def resolve_safetensors_device(map_location: str | torch.device) -> str:
|
||||
"""Resolve a device string for a safetensors load, working around a device-mapping quirk.
|
||||
|
||||
safetensors' load maps the bare string "cuda" to cuda:0 regardless of the current device
|
||||
(unlike torch's .to("cuda"), which honors torch.cuda.current_device()). Under multi-GPU
|
||||
accelerate/FSDP every rank would then load its weights onto GPU 0, OOMing it before sharding.
|
||||
Resolve "cuda" to the concrete current-device index so each rank loads onto its own GPU.
|
||||
"""
|
||||
map_location = str(map_location)
|
||||
if map_location == "cuda" and torch.cuda.is_available():
|
||||
return f"cuda:{torch.cuda.current_device()}"
|
||||
return map_location
|
||||
|
||||
|
||||
def get_safe_dtype(dtype: torch.dtype, device: str | torch.device):
|
||||
"""
|
||||
mps is currently not compatible with float64
|
||||
|
||||
@@ -104,7 +104,6 @@ class MetricsTracker:
|
||||
"episodes",
|
||||
"epochs",
|
||||
"accelerator",
|
||||
"_caller_metrics",
|
||||
]
|
||||
|
||||
def __init__(
|
||||
@@ -130,9 +129,6 @@ class MetricsTracker:
|
||||
self.episodes = self.samples / self._avg_samples_per_ep
|
||||
self.epochs = self.samples / self._num_frames
|
||||
self.accelerator = accelerator
|
||||
# Meter names the caller registered up front. update_metrics() leaves these untouched, so a
|
||||
# policy that echoes e.g. "loss" in its output dict can't clobber the aggregated meter.
|
||||
self._caller_metrics: set[str] = set(self.metrics)
|
||||
|
||||
def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any:
|
||||
if name in self.__dict__:
|
||||
@@ -160,21 +156,6 @@ class MetricsTracker:
|
||||
self.episodes = self.samples / self._avg_samples_per_ep
|
||||
self.epochs = self.samples / self._num_frames
|
||||
|
||||
def update_metrics(self, values: dict[str, Any]) -> None:
|
||||
"""Accumulate a dict of scalar metrics, auto-registering a meter for each new key.
|
||||
|
||||
Non-numeric values and bools are ignored.
|
||||
Caller-registered metrics (those passed to the constructor) are never overridden.
|
||||
"""
|
||||
for name, value in values.items():
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
continue
|
||||
if name in self._caller_metrics:
|
||||
continue
|
||||
if name not in self.metrics:
|
||||
self.metrics[name] = AverageMeter(name, ":.3f", reduction="mean")
|
||||
self.metrics[name].update(float(value))
|
||||
|
||||
def reduce_across_ranks(self) -> None:
|
||||
"""
|
||||
Synchronises the running averages of every metric whose ``reduction`` is not ``"none"``
|
||||
|
||||
@@ -85,7 +85,7 @@ def _spy_responder(captured: list[list[dict[str, Any]]], reply: Any):
|
||||
def test_module1_plan_memory_subtask_smoke(fixture_dataset_root: Path, tmp_path: Path) -> None:
|
||||
vlm = make_canned_responder(
|
||||
{
|
||||
"COMPLETED manipulation events": {
|
||||
"atomic subtasks": {
|
||||
"subtasks": [
|
||||
{"text": "grasp the handle of the sponge", "start": 0.0, "end": 0.4},
|
||||
{"text": "wipe the counter from left to right", "start": 0.4, "end": 0.8},
|
||||
@@ -126,7 +126,7 @@ def test_module1_emit_memory_false_skips_memory_keeps_subtasks_and_plan(
|
||||
leaving subtask + plan generation intact — symmetric to ``emit_plan``."""
|
||||
vlm = make_canned_responder(
|
||||
{
|
||||
"COMPLETED manipulation events": {
|
||||
"atomic subtasks": {
|
||||
"subtasks": [
|
||||
{"text": "grasp the handle of the sponge", "start": 0.0, "end": 0.4},
|
||||
{"text": "wipe the counter from left to right", "start": 0.4, "end": 0.8},
|
||||
@@ -318,7 +318,7 @@ def test_module1_attaches_contact_sheets_to_subtask_prompt(
|
||||
return block.get("text", "")
|
||||
return ""
|
||||
|
||||
subtask_calls = [m for m in captured if "COMPLETED manipulation events" in _prompt_text(m)]
|
||||
subtask_calls = [m for m in captured if "atomic subtasks" in _prompt_text(m)]
|
||||
assert len(subtask_calls) == 1, "expected exactly one subtask-prompt VLM call"
|
||||
content = subtask_calls[0][0]["content"]
|
||||
video_blocks = [b for b in content if isinstance(b, dict) and b.get("type") == "video"]
|
||||
|
||||
@@ -346,10 +346,3 @@ def test_state_not_modified_by_relative_processor(dataset, action_dim):
|
||||
|
||||
result_state = result[TransitionKey.OBSERVATION][OBS_STATE]
|
||||
torch.testing.assert_close(result_state, original_state)
|
||||
|
||||
|
||||
def test_cached_anchor_not_in_config():
|
||||
"""The cached anchor is ephemeral runtime state and must not leak into the config."""
|
||||
step = RelativeActionsProcessorStep(enabled=True)
|
||||
step.set_cached_state(torch.tensor([[1.0, 2.0, 3.0, 4.0]]))
|
||||
assert set(step.get_config()) == {"enabled", "exclude_joints", "action_names"}
|
||||
|
||||
@@ -18,8 +18,6 @@ import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from huggingface_hub.errors import RevisionNotFoundError
|
||||
|
||||
# ``lerobot.scripts.lerobot_annotate`` (and the ``_push_to_hub`` path it
|
||||
# exercises) imports ``lerobot.datasets``, which only ships under the
|
||||
@@ -28,13 +26,11 @@ pytest.importorskip("datasets", reason="datasets is required (install lerobot[da
|
||||
|
||||
|
||||
def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
|
||||
from lerobot.scripts import lerobot_annotate
|
||||
from lerobot.scripts.lerobot_annotate import _push_to_hub
|
||||
|
||||
root = tmp_path / "dataset"
|
||||
(root / "meta").mkdir(parents=True)
|
||||
(root / "meta" / "info.json").write_text(
|
||||
json.dumps({"codebase_version": "v3.0", "fps": 30, "features": {}})
|
||||
)
|
||||
(root / "meta" / "info.json").write_text(json.dumps({"codebase_version": "v3.0"}))
|
||||
|
||||
calls = {}
|
||||
|
||||
@@ -47,6 +43,9 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
|
||||
return SimpleNamespace(oid="abc123")
|
||||
|
||||
def delete_tag(self, repo_id, **kwargs):
|
||||
import requests
|
||||
from huggingface_hub.errors import RevisionNotFoundError
|
||||
|
||||
calls["delete_tag"] = {"repo_id": repo_id, **kwargs}
|
||||
# Simulate the common case: no stale tag to delete.
|
||||
raise RevisionNotFoundError("no such tag", response=requests.Response())
|
||||
@@ -54,12 +53,7 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
|
||||
def create_tag(self, **kwargs):
|
||||
calls["create_tag"] = kwargs
|
||||
|
||||
monkeypatch.setattr(lerobot_annotate, "HfApi", FakeHfApi)
|
||||
|
||||
def fake_card_push(self, **kwargs):
|
||||
calls["card_push"] = {"content": str(self), **kwargs}
|
||||
|
||||
monkeypatch.setattr("huggingface_hub.DatasetCard.push_to_hub", fake_card_push)
|
||||
monkeypatch.setattr("huggingface_hub.HfApi", FakeHfApi)
|
||||
|
||||
cfg = SimpleNamespace(
|
||||
repo_id="source/dataset",
|
||||
@@ -68,7 +62,7 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
|
||||
push_commit_message=None,
|
||||
)
|
||||
|
||||
lerobot_annotate._push_to_hub(root, cfg)
|
||||
_push_to_hub(root, cfg)
|
||||
|
||||
assert calls["create_repo"] == {
|
||||
"repo_id": "annotated/dataset",
|
||||
@@ -77,13 +71,6 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
|
||||
"exist_ok": True,
|
||||
}
|
||||
assert calls["upload_folder"]["repo_id"] == "annotated/dataset"
|
||||
# The source README must not be copied over: its links (e.g. the
|
||||
# visualize badge) point at the source dataset. A card regenerated for
|
||||
# the target repo is pushed instead.
|
||||
assert "README.md" in calls["upload_folder"]["ignore_patterns"]
|
||||
assert calls["card_push"]["repo_id"] == "annotated/dataset"
|
||||
assert "visualize_dataset?path=annotated/dataset" in calls["card_push"]["content"]
|
||||
assert "source/dataset" not in calls["card_push"]["content"]
|
||||
# A stale tag (e.g. from a previous annotation run) is deleted first so
|
||||
# the new tag always points at the upload we just made.
|
||||
assert calls["delete_tag"] == {
|
||||
|
||||
@@ -348,225 +348,3 @@ def test_rollout_context_fields():
|
||||
|
||||
field_names = {f.name for f in dataclasses.fields(RolloutContext)}
|
||||
assert field_names == {"runtime", "hardware", "policy", "processors", "data"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sync engine: relative-action anchoring (drift-free chunk execution)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_REL_ACTION_NAMES = ["j0.pos", "j1.pos", "j2.pos", "gripper.pos"]
|
||||
_REL_ACTION_DIM = len(_REL_ACTION_NAMES)
|
||||
|
||||
|
||||
def _relative_pre_post(exclude_joints=None):
|
||||
"""Pre/post processors wrapping the real relative (caches anchor) and absolute
|
||||
(relative + cached state) steps, mirroring what the sync engine feeds them."""
|
||||
from lerobot.processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
RelativeActionsProcessorStep,
|
||||
TransitionKey,
|
||||
create_transition,
|
||||
)
|
||||
from lerobot.utils.constants import OBS_STATE
|
||||
|
||||
relative_step = RelativeActionsProcessorStep(
|
||||
enabled=True, exclude_joints=exclude_joints or [], action_names=list(_REL_ACTION_NAMES)
|
||||
)
|
||||
absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
|
||||
|
||||
class _Pre:
|
||||
steps = [relative_step]
|
||||
|
||||
def __call__(self, observation):
|
||||
# Run the relative step so it caches the anchor, then pass the batch through.
|
||||
transition = create_transition(observation={OBS_STATE: observation[OBS_STATE]})
|
||||
relative_step(transition)
|
||||
return observation
|
||||
|
||||
def reset(self):
|
||||
pass
|
||||
|
||||
class _Post:
|
||||
def __call__(self, action):
|
||||
transition = create_transition(action=action)
|
||||
return absolute_step(transition)[TransitionKey.ACTION]
|
||||
|
||||
def reset(self):
|
||||
pass
|
||||
|
||||
return _Pre(), _Post(), relative_step
|
||||
|
||||
|
||||
def _fake_relative_policy(chunk_rel, n_action_steps, chunking=True):
|
||||
"""Fake relative-action policy for the sync engine.
|
||||
|
||||
``chunking=True`` buffers a chunk and serves it one action per tick, calling the
|
||||
public ``predict_action_chunk`` only on refill (pi0/fastwam/lingbot). ``False``
|
||||
returns an action directly and never calls it. The engine's anchor probe keys off
|
||||
that public call, so the fake routes through it rather than any private queue.
|
||||
"""
|
||||
from collections import deque
|
||||
|
||||
policy = MagicMock()
|
||||
policy.config.use_amp = False
|
||||
policy.config.action_feature_names = list(_REL_ACTION_NAMES)
|
||||
state = {"predict_calls": 0}
|
||||
queue = deque(maxlen=n_action_steps)
|
||||
|
||||
def predict_action_chunk(_batch=None, **_kwargs):
|
||||
state["predict_calls"] += 1
|
||||
return chunk_rel.unsqueeze(0) # [B=1, n, dim]
|
||||
|
||||
def select_action(_observation):
|
||||
if not chunking:
|
||||
return chunk_rel[0].unsqueeze(0)
|
||||
if len(queue) == 0:
|
||||
actions = policy.predict_action_chunk(_observation)
|
||||
queue.extend(actions.transpose(0, 1)) # [n, 1, dim]
|
||||
return queue.popleft()
|
||||
|
||||
policy.predict_action_chunk.side_effect = predict_action_chunk
|
||||
policy.select_action.side_effect = select_action
|
||||
policy.reset.side_effect = queue.clear
|
||||
policy._predict_state = state
|
||||
return policy
|
||||
|
||||
|
||||
def _build_sync_engine(policy, pre, post):
|
||||
from lerobot.rollout import SyncInferenceEngine
|
||||
|
||||
return SyncInferenceEngine(
|
||||
policy=policy,
|
||||
preprocessor=pre,
|
||||
postprocessor=post,
|
||||
dataset_features={"action": {"names": list(_REL_ACTION_NAMES)}},
|
||||
ordered_action_keys=list(_REL_ACTION_NAMES),
|
||||
task="test",
|
||||
device="cpu",
|
||||
robot_type="mock",
|
||||
)
|
||||
|
||||
|
||||
def _obs_frame(state_values):
|
||||
import numpy as np
|
||||
|
||||
return {"observation.state": np.asarray(state_values, dtype=np.float32)}
|
||||
|
||||
|
||||
def test_sync_relative_holds_anchor_across_chunk():
|
||||
"""Every action popped within a chunk must anchor to the tick-0 state (no drift)."""
|
||||
n = 4
|
||||
# A distinct relative offset per chunk step so a wrong anchor would be visible.
|
||||
chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.1 * (i + 1)) for i in range(n)])
|
||||
pre, post, relative_step = _relative_pre_post()
|
||||
policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
|
||||
engine = _build_sync_engine(policy, pre, post)
|
||||
|
||||
assert engine._relative_step is relative_step # introspection wired the step
|
||||
|
||||
s0 = [1.0, 2.0, 3.0, 4.0]
|
||||
outputs = []
|
||||
for tick in range(n):
|
||||
# Feed a *different* state each tick; a drifting anchor would use it.
|
||||
state = [v + tick for v in s0]
|
||||
outputs.append(engine.get_action(_obs_frame(state)))
|
||||
|
||||
# Exactly one chunk was predicted across the n ticks.
|
||||
assert policy._predict_state["predict_calls"] == 1
|
||||
for tick in range(n):
|
||||
expected = torch.tensor(s0) + chunk_rel[tick]
|
||||
torch.testing.assert_close(outputs[tick], expected)
|
||||
|
||||
# Next tick empties the queue -> fresh chunk -> anchor advances to the new state.
|
||||
s_next = [10.0, 20.0, 30.0, 40.0]
|
||||
out = engine.get_action(_obs_frame(s_next))
|
||||
assert policy._predict_state["predict_calls"] == 2
|
||||
torch.testing.assert_close(out, torch.tensor(s_next) + chunk_rel[0])
|
||||
# The anchor now reflects the fresh-chunk state, not the held one.
|
||||
torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_next]))
|
||||
|
||||
|
||||
def test_sync_relative_reset_reanchors_new_episode():
|
||||
"""After ``reset()`` the first tick of the next episode anchors to the new state."""
|
||||
n = 3
|
||||
chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)])
|
||||
pre, post, relative_step = _relative_pre_post()
|
||||
policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
|
||||
engine = _build_sync_engine(policy, pre, post)
|
||||
|
||||
# Episode 1: one tick anchors to s0 and leaves cached actions in the queue.
|
||||
engine.get_action(_obs_frame([1.0, 1.0, 1.0, 1.0]))
|
||||
assert policy._predict_state["predict_calls"] == 1
|
||||
|
||||
engine.reset() # clears the queue and the per-episode chunk flags
|
||||
|
||||
# Episode 2: a fresh state must produce a fresh chunk anchored to that state,
|
||||
# not carry over the previous episode's anchor.
|
||||
s_new = [7.0, 8.0, 9.0, 10.0]
|
||||
out = engine.get_action(_obs_frame(s_new))
|
||||
assert policy._predict_state["predict_calls"] == 2
|
||||
torch.testing.assert_close(out, torch.tensor(s_new) + chunk_rel[0])
|
||||
torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_new]))
|
||||
|
||||
|
||||
def test_sync_relative_non_chunking_policy_refreshes_every_tick():
|
||||
"""A policy that never calls ``predict_action_chunk`` must not freeze the anchor."""
|
||||
n = 3
|
||||
chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.5) for _ in range(n)])
|
||||
pre, post, _ = _relative_pre_post()
|
||||
policy = _fake_relative_policy(chunk_rel, n_action_steps=n, chunking=False)
|
||||
engine = _build_sync_engine(policy, pre, post)
|
||||
|
||||
s0 = [1.0, 1.0, 1.0, 1.0]
|
||||
for tick in range(3):
|
||||
state = [v + tick for v in s0]
|
||||
out = engine.get_action(_obs_frame(state))
|
||||
# Anchor must track the current state every tick (no chunk => no hold).
|
||||
torch.testing.assert_close(out, torch.tensor(state) + chunk_rel[0])
|
||||
assert policy._predict_state["predict_calls"] == 0
|
||||
|
||||
|
||||
def test_sync_engine_no_relative_step_is_none():
|
||||
"""Without an enabled relative step, the engine takes the plain select_action path."""
|
||||
policy = MagicMock()
|
||||
policy.config.use_amp = False
|
||||
engine = _build_sync_engine(policy, MagicMock(steps=[]), MagicMock())
|
||||
assert engine._relative_step is None
|
||||
|
||||
|
||||
def test_sync_relative_exclude_joints_stay_absolute():
|
||||
"""With ``exclude_joints``, excluded dims pass through absolute while the relative
|
||||
dims still hold the tick-0 anchor across the chunk."""
|
||||
n = 4
|
||||
# Distinct offset per step *and* per dim so a wrong anchor or a wrong mask shows up.
|
||||
chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.1 * (i + 1)) for i in range(n)])
|
||||
pre, post, relative_step = _relative_pre_post(exclude_joints=["gripper"])
|
||||
policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
|
||||
engine = _build_sync_engine(policy, pre, post)
|
||||
|
||||
# gripper (last dim) is kept absolute; j0..j2 are relative.
|
||||
mask = torch.tensor([1.0, 1.0, 1.0, 0.0])
|
||||
s0 = [1.0, 2.0, 3.0, 4.0]
|
||||
outputs = []
|
||||
for tick in range(n):
|
||||
state = [v + tick for v in s0] # moving state; a drifting anchor would use it
|
||||
outputs.append(engine.get_action(_obs_frame(state)))
|
||||
|
||||
assert policy._predict_state["predict_calls"] == 1 # one chunk held across n ticks
|
||||
for tick in range(n):
|
||||
# relative dims: anchor held at s0; excluded gripper dim: raw predicted value.
|
||||
expected = chunk_rel[tick] + torch.tensor(s0) * mask
|
||||
torch.testing.assert_close(outputs[tick], expected)
|
||||
|
||||
|
||||
def test_sync_relative_stop_restores_policy_method():
|
||||
"""``stop()`` un-patches the probe so the policy object isn't permanently modified."""
|
||||
n = 3
|
||||
chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)])
|
||||
pre, post, _ = _relative_pre_post()
|
||||
policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
|
||||
original = policy.predict_action_chunk
|
||||
engine = _build_sync_engine(policy, pre, post)
|
||||
assert policy.predict_action_chunk is not original # probe installed
|
||||
engine.stop()
|
||||
assert policy.predict_action_chunk is original # restored
|
||||
|
||||
@@ -233,37 +233,3 @@ def test_metrics_tracker_reduce_across_ranks_invokes_reduce():
|
||||
# accumulate against the cluster view rather than the stale per-rank sum.
|
||||
meter = tracker.update_s
|
||||
assert meter.sum / meter.count == pytest.approx(meter.avg)
|
||||
|
||||
|
||||
def test_metrics_tracker_update_metrics_registers_and_averages():
|
||||
tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics={})
|
||||
tracker.update_metrics({"latent_loss": 0.2, "action_loss": 0.4})
|
||||
tracker.update_metrics({"latent_loss": 0.4, "action_loss": 0.6})
|
||||
|
||||
# New keys are auto-registered as mean-reduced meters and averaged over the window.
|
||||
assert tracker.metrics["latent_loss"].reduction == "mean"
|
||||
assert tracker.metrics["latent_loss"].avg == pytest.approx(0.3)
|
||||
assert tracker.metrics["action_loss"].avg == pytest.approx(0.5)
|
||||
assert tracker.to_dict()["latent_loss"] == pytest.approx(0.3)
|
||||
|
||||
|
||||
def test_metrics_tracker_update_metrics_skips_non_numeric():
|
||||
tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics={})
|
||||
tracker.update_metrics({"loss": 0.5, "head_mode": "sparse", "enabled": True})
|
||||
|
||||
# strings and bools ignored
|
||||
assert "loss" in tracker.metrics
|
||||
assert "head_mode" not in tracker.metrics
|
||||
assert "enabled" not in tracker.metrics
|
||||
|
||||
|
||||
def test_metrics_tracker_update_metrics_does_not_override_caller_meter():
|
||||
# A policy that echoes "loss" in its output dict must not overwrite the caller-owned,
|
||||
# already-aggregated loss meter.
|
||||
metrics = {"loss": AverageMeter("loss", ":.3f", reduction="mean")}
|
||||
tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=metrics)
|
||||
tracker.loss = 1.0 # caller-set optimized loss
|
||||
tracker.update_metrics({"loss": 99.0, "latent_loss": 0.2})
|
||||
|
||||
assert tracker.metrics["loss"].avg == pytest.approx(1.0) # snapshot ignored
|
||||
assert tracker.metrics["latent_loss"].avg == pytest.approx(0.2)
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
version = 1
|
||||
revision = 3
|
||||
revision = 2
|
||||
requires-python = ">=3.12"
|
||||
resolution-markers = [
|
||||
"(python_full_version >= '3.15' and platform_machine == 'AMD64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux')",
|
||||
@@ -2127,16 +2127,16 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "hydra-core"
|
||||
version = "1.3.4"
|
||||
version = "1.3.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "antlr4-python3-runtime", marker = "sys_platform == 'linux'" },
|
||||
{ name = "omegaconf", marker = "sys_platform == 'linux'" },
|
||||
{ name = "packaging", marker = "sys_platform == 'linux'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/10/dd/220f0e91743136725352497e98540772a01fc7c3ab96ff16c3c74424e984/hydra_core-1.3.4.tar.gz", hash = "sha256:ad0f7b05a0242255a8984d5a4ed2f6847f7b783ed727368a2c0155ec52d6c34c", size = 3263348, upload-time = "2026-07-04T16:25:38.891Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/0d/0b/7c0d941311aadc6479ec01767edba9c8a07db1452685de3567ed3058d0c9/hydra_core-1.3.3.tar.gz", hash = "sha256:b7477ee21f08b62f71bf0126d44695c048dc7e9c0cc79e2d593b707cb1e44048", size = 3262532, upload-time = "2026-06-11T05:54:26.835Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/ed/cd/a568610bafe991fdd3f628fb606316b3b2be52ded019284e895d9beb3a1e/hydra_core-1.3.4-py3-none-any.whl", hash = "sha256:e58683692904a09f1fdfffa1a9b86bfd94e215b59f1ee17e7cd7d92738090d33", size = 155478, upload-time = "2026-07-04T16:25:37.291Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e5/57/4e39f85347f77144d2ad12e87d5df8fb8f17023f9bd9e8c6e903a128382c/hydra_core-1.3.3-py3-none-any.whl", hash = "sha256:cf349fc393f486f250e5825592c3d0a50c0af3effd726cf8dd5b637a7cb464e3", size = 154706, upload-time = "2026-06-11T05:54:24.917Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2823,7 +2823,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "lerobot"
|
||||
version = "0.6.1"
|
||||
version = "0.5.2"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "cmake" },
|
||||
@@ -6934,15 +6934,15 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "uvicorn"
|
||||
version = "0.50.0"
|
||||
version = "0.49.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "click" },
|
||||
{ name = "h11" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/2e/41/06cce5dbb9f77591512957710ac709e60b12e6216a2f2d0d607fd49706e8/uvicorn-0.50.0.tar.gz", hash = "sha256:0c92e1bc2259cb7faa4fcef774a5966588f2e88542744550b66799fba10b76f1", size = 93257, upload-time = "2026-07-04T05:03:26.33Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/c4/1f/fa18009dea8469069cca78a4e877a008ab78f08b064bfc9ab891579077ff/uvicorn-0.49.0.tar.gz", hash = "sha256:ebf4271aa580d9de97f93192d4595176df6e91f9aae919ca73e4fc07df1e66a3", size = 91284, upload-time = "2026-06-03T22:01:30.448Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a0/3a/eb70620ca2bf8213603d5c731460687c49fee38b0072f0b4a637781f0a53/uvicorn-0.50.0-py3-none-any.whl", hash = "sha256:05f0eb19edf38208f79f43df8a63081b48df31b0cd1e5997be957a4dc97d1b19", size = 72716, upload-time = "2026-07-04T05:03:24.848Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/88/fa/e1388bbcf24ef3274f45c0c1c7b501fd14971037c1b6ee23610553307497/uvicorn-0.49.0-py3-none-any.whl", hash = "sha256:ba3d14c3ee7e41c6c654c46c9eb489d33213cdd30aa1696eab1374337c13f68f", size = 71376, upload-time = "2026-06-03T22:01:29.037Z" },
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
|
||||
Reference in New Issue
Block a user