diff --git a/docs/source/pi052.mdx b/docs/source/pi052.mdx index 73ccd6236..1d5accc35 100644 --- a/docs/source/pi052.mdx +++ b/docs/source/pi052.mdx @@ -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 diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py index 1d41a1f7c..77bb8f6b3 100644 --- a/src/lerobot/policies/pi052/configuration_pi052.py +++ b/src/lerobot/policies/pi052/configuration_pi052.py @@ -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