action dataloader: episode-shuffle stream (fix DROID grad-norm instability)

[fwd4]

Problem

The DROID action SFT dataloader trained with an unstable, slow-settling grad-norm (and a noisy action-loss plateau) vs the internal reference. Root cause: the DROID action dataset is map-style and — unlike the iterable vision SFTDataset, which self-shuffles — does not shuffle, and RankPartitionedDataLoader wraps it in a DataLoader with no shuffle, i.e. a SequentialSampler. Every rank then iterates the same consecutive, overlapping windows, so the all-reduced global batch is effectively ~1 episode → high gradient variance.

(Forward + gradients were verified numerically equivalent to the internal model on identical input, so this was a data-path issue, not the model/loss/optimizer.)

Fix

ActionIterableShuffleDataset (iterable_shuffle=True): an IterableDataset view of the map-style dataset that streams rank × worker-sharded, episode-order-shuffled, sequential-within-episode — decorrelated batches with sequential reads (preserves I/O locality + copy-on-write; a plain shuffle=True/RandomSampler instead does random-access I/O → ~11 min/iter and OOM from broken COW). Mirrors the internal iterable dataset's per-worker episode assignment.

• Adds DROIDLeRobotDataset.get_shuffle_blocks() (per-episode/segment flat-index blocks the iterable streams).
• No DataLoader/sampler change needed — IterableDataset is handled natively (sampler=None).

Validation (8192 global batch)

iter	this fix	internal ref	no-shuffle
100	grad-norm 2.9	4.7	21
450	grad-norm 1.7	1.9	—

Per-component action loss converges to ~0.0055 (matches internal ~0.005; the no-shuffle run plateaued noisily at 0.03–0.07). Builds on #24 (recipe + FusedAdam optimizer).

🤖 Generated with Claude Code

---

Added commits (recipe correctness)

• mode="policy" default — DROIDLeRobotDataset defaulted to mode="joint" (random forward_dynamics/inverse_dynamics/policy per sample), so the policy recipe was silently training multi-task. inverse_dynamics zeros the vision loss and forward_dynamics zeros the action loss, diluting each per-task loss by ~1/3 vs the policy-only internal run. Now defaults to policy (matching i4's DROIDLeRobotDataset); mode is also threaded through get_action_droid_sft_dataset.
• max_num_tokens_after_packing=-1 — uncaps the packed-sequence length (NANO default 45056) to match the internal droid_lerobot_8b run, so the full vision sequence is processed per step. Does not change the per-token loss; widens the effective vision context per step.

[mli0603]

LGTM

[lfengad]

overall LGTM

[ychao-nvidia]

Is this change intended? Our internal ablation showed that fixing lr to 2.0e-4 is a key to high policy success rate.