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

fix(pi052): use base learning rate for lm head

Browse Source
This commit is contained in:
Pepijn
2026-07-15 15:06:22 +02:00
parent 6a89c7be45
commit 5242e9195c
2 changed files with 12 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/pi052.mdx
+10 -10
View File
@@ -107,16 +107,16 @@ the expected prompt, text target, and action endpoints before scaling up.
### Main training controls
| Option | Default | Purpose |
| -------------------------------- | -------------------------: | ------------------------------------------------------------------ |
| `policy.recipe_path` | `recipes/subtask_mem.yaml` | Selects the language/action objective mixture |
| `policy.text_loss_weight` | `1.0` | Language-head cross-entropy weight; `0` disables text training |
| `policy.flow_loss_weight` | `10.0` | Continuous action flow-loss weight |
| `policy.enable_fast_action_loss` | `true` | Adds discrete FAST action-token supervision |
| `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.lm_head_lr_scale` | `5.0` | Gives the sparsely supervised language head a larger learning rate |
| Option | Default | Purpose |
| -------------------------------- | -------------------------: | -------------------------------------------------------------- |
| `policy.recipe_path` | `recipes/subtask_mem.yaml` | Selects the language/action objective mixture |
| `policy.text_loss_weight` | `1.0` | Language-head cross-entropy weight; `0` disables text training |
| `policy.flow_loss_weight` | `10.0` | Continuous action flow-loss weight |
| `policy.enable_fast_action_loss` | `true` | Adds discrete FAST action-token supervision |
| `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.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
src/lerobot/policies/pi052/configuration_pi052.py
+2 -2
View File
@@ -120,8 +120,8 @@ class PI052Config(PI05Config):
"""If True, route every transformer layer through the KI
attention path that blocks action→VLM gradient flow on K/V."""
# Boost sparse text-head updates while retaining PI0.5's optimizer schedule.
lm_head_lr_scale: float = 5.0
# Scale language-head updates relative to the base optimizer schedule.
lm_head_lr_scale: float = 1.0
# Scale pretrained backbone and new action-expert groups independently; 1.0 preserves legacy behavior.
backbone_lr_scale: float = 1.0