Co-authored-by: Remi <re.cadene@gmail.com>

Co-authored-by: Remi <re.cadene@gmail.com>

Co-authored-by: Remi <re.cadene@gmail.com>

For 32-dim sparse action layouts where only ~10 dims carry real signal
(airoa HSR with pi05 baseline multi-robot layout), the other 22 zero-padded
dims dilute the flow matching loss and cause active dims (especially gripper)
to under-learn.

Add optional config field `action_loss_exclude_dims: list[int] | None`:
- When None (default): no change, backward compatible
- When specified: mask listed dims from flow loss sum AND adjust denominator
  to use only included dims, so active dim signal is preserved

Related: icra_2026_ramen Issue huggingface#125 (Run 61)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

… loss

Two optional loss-shape modifications for PI0.5 flow matching, motivated
by the gripper-underlearning issue observed in icra_2026_ramen Run 60/61
(Issue huggingface#125 on the icra side).

1. flow_loss_type: "mse" (default) | "smooth_l1"
   Swaps F.mse_loss for F.smooth_l1_loss (Huber) with configurable beta.
   Empirical precedent: Kim, Finn et al. 2025 (arxiv 2502.19645) found
   L1/Huber regression beat classification heads on OpenVLA LIBERO
   fine-tuning — directly contradicting the DAFD (gripper=CE) approach
   tried in Run 60 which hurt accuracy.

2. use_min_snr_weighting + min_snr_gamma
   Min-SNR-γ (Hang et al. 2023, arxiv 2303.09556) v-prediction variant
   applied to flow matching via the diffusion↔FM weighting equivalence
   (diffusionflow.github.io). Per-sample weight:
       w(t) = min(SNR(t), γ) / (SNR(t) + 1)
       SNR(t) = ((1-t)/t)^2   (linear FM schedule)
   Peaks at SNR=γ (≈ t=0.31 with γ=5), downweighting both very-low and
   very-high noise timesteps. Default γ=5.0 from original paper.

Both defaults preserve current behavior (mse, no weighting).

Logging:
- flow_loss_raw: pre-weighting loss (for cross-run comparison)
- flow_loss: final (post-weighting) loss
- min_snr_weight_mean: diagnostic for the Min-SNR weight batch average

Tests (tests/policies/pi0_pi05/test_pi05_loss_variants.py):
- Config validation for new fields + rejection of invalid values
- Min-SNR formula numerical correctness (scalar + vectorized)
- SmoothL1 produces smaller loss than MSE on outlier inputs

Related: icra_2026_ramen Issue huggingface#125 (Run 62)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…guard

Follow-up on self-review: dropping the SmoothL1 variant from this PR
because the cited evidence (Kim-Finn OFT, arxiv 2502.19645) does not
actually ablate L1 vs MSE — it compares L1 vs discrete tokens vs
diffusion head in a single-step regression setting, not flow matching
velocity. The direct empirical case for swapping MSE → SmoothL1 in FM
VLA is therefore weaker than claimed in the original PR body.

Also: SmoothL1 with beta=1.0 behaves as MSE for |error| < 1, and FM
velocity residuals on normalized actions are typically below 1, so the
change would be near no-op most of the time.

Keeping Min-SNR only makes this Run 62 change a single-variable
intervention, which preserves causal attribution vs Run 61 (mask only).

Additions:
- Terminal-SNR guard: follow original Hang et al. repo
  (guided_diffusion/gaussian_diffusion.py:895) and set weight=1.0 when
  SNR==0. With the eps=1e-6 t-clamp this does not fire in practice,
  but it guards the edge if a caller passes t=1.0 exactly.
- Test for the guard behavior.

Related: icra_2026_ramen Issue huggingface#125

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

The lerobot wandb wrapper emits a warning and drops any loss_dict value
whose type is not a scalar, so tensors / lists of shape (action_dim,)
are silently ignored. Previously "loss_per_dim", "loss_per_dim_weighted",
and "smoothness_per_dim" were all dropped, making per-dim analysis
impossible without modifying the wrapper.

Expand each tensor into individual scalar entries:
  loss_per_dim           -> loss_per_dim_00, loss_per_dim_01, ..., loss_per_dim_31
  loss_per_dim_weighted  -> loss_per_dim_weighted_{NN}
  smoothness_per_dim     -> smoothness_per_dim_{NN}

This is required for icra_2026_ramen Issue huggingface#125 Run 62 monitoring where
gripper (dim 6 in 32-dim sparse layout) per-dim loss is the go/no-go
signal at 5-10k steps. Applies to Run 63+ as well.

Related: icra_2026_ramen Issue huggingface#125

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>