admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-26 03:06:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(pi052): add training-time RTC

Browse Source
This commit is contained in:
Pepijn
2026-07-16 10:38:28 +02:00
parent 5b8e6ffe8e
commit 18e02ded4f
10 changed files with 396 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/inference.mdx
+7 -6
View File
@@ -274,12 +274,13 @@ lerobot-rollout \
--device=cuda
```
| Flag | Description |
| ------------------------------------------- | -------------------------------------------------------------- |
| `--inference.rtc.execution_horizon` | Steps to blend with previous chunk (default: varies by policy) |
| `--inference.rtc.max_guidance_weight` | Consistency enforcement strength (default: varies by policy) |
| `--inference.rtc.prefix_attention_schedule` | Blend schedule: `LINEAR`, `EXP`, `ONES`, `ZEROS` |
| `--inference.queue_threshold` | Max queue size before backpressure (default: 30) |
| Flag | Description |
| ------------------------------------------- | ------------------------------------------------------------------------------- |
| `--inference.rtc.execution_horizon` | Steps to blend with previous chunk (default: varies by policy) |
| `--inference.rtc.mode` | `guided` (default) or trained-prefix `trained` for compatible Pi052 checkpoints |
| `--inference.rtc.max_guidance_weight` | Consistency enforcement strength (default: varies by policy) |
| `--inference.rtc.prefix_attention_schedule` | Blend schedule: `LINEAR`, `EXP`, `ONES`, `ZEROS` |
| `--inference.queue_threshold` | Max queue size before backpressure (default: 30) |
See the [Real-Time Chunking](./rtc) guide for details on tuning RTC parameters.
docs/source/pi052.mdx
+59
View File
@@ -116,12 +116,71 @@ the expected prompt, text target, and action endpoints before scaling up.
| `policy.fast_action_loss_weight` | `1.0` | FAST cross-entropy weight |
| `policy.knowledge_insulation` | `true` | Blocks action-loss gradients through the VLM K/V path |
| `policy.flow_num_repeats` | `5` | Reuses one VLM prefix for independent denoising targets |
| `policy.rtc_training_max_delay` | `0` | Maximum clean-prefix delay; `0` disables training-time RTC |
| `policy.lm_head_lr_scale` | `1.0` | Scales language-head learning rate; `1.0` uses the base rate |
The loss weights are starting points, not dataset-independent constants. Track
flow loss and text/FAST losses separately, and inspect generated subtasks rather
than selecting a checkpoint from total loss alone.
### Training-time RTC
Pi052 optionally supports training-time action conditioning from
[Training-Time Action Conditioning for Efficient Real-Time Chunking](https://arxiv.org/abs/2512.05964).
It simulates inference latency by sampling a clean action prefix for every flow
draw, passing a per-action flow timestep to the action expert, and computing the
flow loss only on the remaining postfix. The default value of `0` leaves the
standard Pi052 objective unchanged.
```bash
lerobot-train \
--dataset.repo_id=${HF_USER}/my_language_annotated_dataset \
--policy.type=pi052 \
--policy.pretrained_path=lerobot/pi05_base \
--policy.recipe_path=recipes/subtask_mem.yaml \
--policy.rtc_training_max_delay=10 \
--policy.dtype=bfloat16 \
--policy.device=cuda \
--batch_size=8 \
--steps=30000 \
--output_dir=outputs/pi052_rtc \
--job_name=pi052_rtc
```
`rtc_training_max_delay` is measured in controller steps and must be smaller
than `chunk_size`. Choose it to cover the largest inference latency expected at
deployment: at 50 Hz, for example, 10 steps correspond to 200 ms. A delay of
zero is included in the uniform sampling distribution, so the checkpoint also
continues to receive ordinary flow-matching examples. Set rollout's
`inference.rtc.execution_horizon` to at least this maximum so the previous
chunk cache retains enough actions to construct every supported prefix.
Run the resulting checkpoint with the asynchronous `lerobot-rollout` backend
and select the trained-prefix path explicitly:
```bash
lerobot-rollout \
--strategy.type=base \
--policy.path=outputs/pi052_rtc/checkpoints/last/pretrained_model \
--inference.type=rtc \
--inference.rtc.mode=trained \
--inference.rtc.execution_horizon=10 \
--robot.type=so100_follower \
--robot.port=/dev/ttyACM0 \
--task="pick up the cube" \
--fps=50 \
--device=cuda
```
The rollout engine measures latency continuously, carries the still-unexecuted
actions from the previous chunk into the next prediction, and discards the
prefix that elapsed during inference. If the measured delay exceeds the
checkpoint's `rtc_training_max_delay`, rollout stops with an explicit error
instead of silently extrapolating beyond the training distribution. Use
`--inference.rtc.mode=guided` for the original Jacobian-guided RTC path; it does
not require a training-time RTC checkpoint but adds backward-pass work during
denoising.
### Dataset-specific FAST tokenizer
The universal FAST tokenizer works out of the box. For a large or
docs/source/rtc.mdx
+27 -1
View File
@@ -1,6 +1,6 @@
# Real-Time Chunking (RTC)
Real-Time Chunking (RTC) is an inference-time method that allows large, flow-matching based robotic policies, such as [Pi0](./pi0), [Pi0.5](./pi05), and [SmolVLA](./smolvla), to produce smooth, continuous, and reactive motion despite having high inference latency.
Real-Time Chunking (RTC) allows large, flow-matching based robotic policies, such as [Pi0](./pi0), [Pi0.5](./pi05), and [SmolVLA](./smolvla), to produce smooth, continuous, and reactive motion despite having high inference latency. LeRobot provides the original inference-time guided mode and, for compatible Pi052 checkpoints, training-time action conditioning with cheap hard-prefix inference.
These policies generate chunks of future actions (e.g., 50 steps at a time) instead of single actions.
Because the models are large, producing each chunk takes longer than the time it takes the robot to execute it.
@@ -92,6 +92,11 @@ for step in range(num_steps):
`RTCConfig` has the following parameters to tune:
**`mode`** selects the action-prefix conditioning method:
- `guided` (default) applies the original Jacobian guidance during denoising and works with ordinary flow-matching checkpoints.
- `trained` hard-inpaints the previous chunk's prefix with per-action flow timesteps. It currently requires a Pi052 checkpoint trained with `policy.rtc_training_max_delay > 0` and avoids the guidance backward pass.
**`execution_horizon`**: How many timesteps from the previous chunk to maintain consistency with. Higher values mean smoother transitions but potentially less reactivity.
Typical values: 8-12 steps
@@ -141,6 +146,7 @@ lerobot-rollout \
--strategy.type=base \
--policy.path=${HF_USERNAME}/policy_repo_id \
--inference.type=rtc \
--inference.rtc.mode=guided \
--inference.rtc.execution_horizon=10 \
--inference.rtc.max_guidance_weight=10.0 \
--robot.type=so100_follower \
@@ -151,6 +157,24 @@ lerobot-rollout \
--device=cuda
```
For a training-time RTC Pi052 checkpoint, change the mode to `trained`. The
checkpoint records its maximum supported delay, and rollout validates measured
latency against it:
```bash
lerobot-rollout \
--strategy.type=base \
--policy.path=${HF_USERNAME}/pi052_training_rtc \
--inference.type=rtc \
--inference.rtc.mode=trained \
--inference.rtc.execution_horizon=10 \
--robot.type=so100_follower \
--robot.port=/dev/tty.usbmodem58FA0834591 \
--task="Move green small object into the purple platform" \
--duration=120 \
--device=cuda
```
## How It Differs from the Async Inference in LeRobot
Both RTC and [async inference](./async) improve real-time robot control, but they solve different problems.
@@ -189,3 +213,5 @@ See `examples/rtc/eval_dataset.py` for a complete example of offline RTC visuali
- [Smooth-As-Butter Robot Policies](https://alexander-soare.github.io/robotics/2025/08/05/smooth-as-butter-robot-policies.html) - Excellent technical explanation with real robot results
- [Physical Intelligence - Real-Time Chunking](https://www.physicalintelligence.company/research/real_time_chunking) - Original paper and research
- [Kinetix RTC Implementation](https://github.com/Physical-Intelligence/real-time-chunking-kinetix) - Reference implementation from Physical Intelligence
- [Training-Time Action Conditioning](https://arxiv.org/abs/2512.05964) - Efficient RTC with clean-prefix conditioning during training
- [RLDX-1](https://github.com/RLWRLD/RLDX-1) - PyTorch reference used for the training-time RTC integration