diff --git a/docs/source/pi052.mdx b/docs/source/pi052.mdx index 43fd8c918..73ccd6236 100644 --- a/docs/source/pi052.mdx +++ b/docs/source/pi052.mdx @@ -111,7 +111,7 @@ the expected prompt, text target, and action endpoints before scaling up. | -------------------------------- | -------------------------: | ------------------------------------------------------------------ | | `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` | `5.0` | Continuous action flow-loss weight | +| `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 | diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py index 6442c33d1..1d41a1f7c 100644 --- a/src/lerobot/policies/pi052/configuration_pi052.py +++ b/src/lerobot/policies/pi052/configuration_pi052.py @@ -48,11 +48,8 @@ class PI052Config(PI05Config): """Weight on the LM-head cross-entropy term. Set to ``0`` to disable text training entirely (reverts to flow-only / π0.5 behaviour).""" - flow_loss_weight: float = 5.0 - """Weight on the action-expert flow-matching term. ``5.0`` — a milder - flow:text split than the paper's α=10, since the flow-only - ``low_level`` recipe already gives the action expert frequent - gradient. Lower it further if the LM head still underfits.""" + flow_loss_weight: float = 10.0 + """Weight on action-expert flow matching relative to text supervision.""" # Backbone training --------------------------------------------------- unfreeze_lm_head: bool = True