add experiment_cfg/train_config.yaml

experiment_cfg/train_config.yaml

@@ -0,0 +1,111 @@
+# GR00T N1.7 CubeStack 3-CLIENT FL arm — the GR00T twin of
+# smolvla_cubestack_fl_3client.yaml / pi05_cubestack_fl_3client.yaml.
+#
+# SAME colour-pair non-IID partition as those two arms (each client owns one
+# colour pair, both stacking directions, on its OWN robot, blind to the 3rd
+# colour):
+#   c0 (harry,  bare Cube_*)  Green/Orange -> blind to Blue
+#   c1 (zhekai, _manuka)      Green/Blue   -> blind to Orange
+#   c2 (kevin,  _kevin)       Orange/Blue  -> blind to Green
+# Expressed with the `explicit_tasks` scheme (each CubeStack dataset carries
+# exactly one distinct task string), so the split is DECLARED rather than
+# derived from task_round_robin's seeded shuffle — it must reproduce the other
+# two policies' split exactly.
+#
+# SCOPE: action-head-only (~1.62B federated) — backbone frozen
+# (tune_llm/tune_visual false). This is NVIDIA's own new-embodiment recipe
+# (train projector + diffusion head) and matches every GR00T row in the paper.
+# `tune_projector` stays ON by default and MUST NOT be disabled: the embodiment
+# projector slot ships at random init for any non-posttrained embodiment.
+#
+# DATA: the Cube_* repos are LeRobot v3.0, which GR00T's loader cannot read
+# (it requires v2.x meta/episodes.jsonl + per-episode parquet/mp4). These paths
+# point at the v2.1 conversion produced by
+# `scripts/convert_cubestack_v30_to_v21.py` (see slurm/convert_cubestack_groot.sh).
+#
+# RUNTIME: flower, NOT LocalRunner — multi-client GR00T on one node OOMs the
+# 115GB host at the client-to-client transition (measured, job 5473891).
+name: groot_cubestack_fl3_fedavg_headonly
+seed: 42
+device: auto
+
+backend:
+  name: groot
+  pretrained_path: nvidia/GR00T-N1.7-3B
+  optim:
+    name: adamw
+    lr: 1.0e-4
+    weight_decay: 1.0e-5
+    betas: [0.9, 0.95]
+    scheduler: cosine
+    warmup_steps: 1000
+    decay_steps: 12500        # per-client budget: 50 rounds x 250 local steps
+    decay_lr: 1.0e-5
+  extra:
+    embodiment_tag: NEW_EMBODIMENT   # SO-101; GR00T ships no config for it -> declared below
+    modality_config:
+      video:
+        delta_indices: [0]
+        modality_keys: [front, wrist]
+      state:
+        delta_indices: [0]
+        modality_keys: [single_arm, gripper]
+      action:
+        delta_indices: 16              # shorthand for list(range(16)) = the action horizon
+        modality_keys: [single_arm, gripper]
+      language:
+        delta_indices: [0]
+        modality_keys: [annotation.human.task_description]
+    state_dropout_prob: 0.2
+    dtype: bfloat16
+    gradient_checkpointing: true
+    local_files_only: true
+    tune_llm: false                    # action-head-only: backbone frozen
+    tune_visual: false
+
+data:
+  name: groot
+  batch_size: 32
+  num_workers: 8                       # MEASURED: each spawn worker ~4GB baseline (fresh torch+gr00t import); 16 hit the 115GB cap at 87GB
+  extra:
+    video_backend: pyav
+    dataloader_mp_context: spawn       # GR00T + CUDA: fork segfaults
+    episode_sampling_rate: 1.0
+    shard_size: 32
+    lazy_episode_datapoints: true      # sampling-PRESERVING speed fix (global shuffle intact)
+    episode_cache_size: 2      # ~230MB/episode at 256-edge
+    prefetch_factor: 2
+    # Order matters: `explicit_tasks` maps task name -> client, but list_episodes
+    # walks these paths in order to build the global episode index.
+    dataset_paths:
+      - /lus/lfs1aip2/projects/u6nf/zhekai.u6nf/groot-cubestack-data-256/Cube_Green_on_Orange
+      - /lus/lfs1aip2/projects/u6nf/zhekai.u6nf/groot-cubestack-data-256/Cube_Orange_on_Green
+      - /lus/lfs1aip2/projects/u6nf/zhekai.u6nf/groot-cubestack-data-256/Cube_Green_on_Blue_manuka
+      - /lus/lfs1aip2/projects/u6nf/zhekai.u6nf/groot-cubestack-data-256/Cube_Blue_on_Green_manuka
+      - /lus/lfs1aip2/projects/u6nf/zhekai.u6nf/groot-cubestack-data-256/Cube_Orange_on_Blue_kevin
+      - /lus/lfs1aip2/projects/u6nf/zhekai.u6nf/groot-cubestack-data-256/Cube_Blue_on_Orange_kevin
+  partitioning:
+    scheme: explicit_tasks
+    num_clients: 3
+    seed: 42
+  # NOTE kevin's two prompts are phrased differently from harry's/manuka's
+  # ("Put orange cube on the blue cube." vs "Put the ... on top of the ...").
+  # Kept VERBATIM: the pi0.5 and SmolVLA CubeStack arms trained on these exact
+  # strings, and rewording here would break comparability with them.
+  client_task_names:
+    - ["Put the green cube on top of the orange cube.", "Put the orange cube on top of the green cube."]
+    - ["Put the green cube on top of the blue cube.", "Put the blue cube on top of the green cube."]
+    - ["Put orange cube on the blue cube.", "Put blue cube on the orange cube."]
+
+fl:
+  strategy:
+    name: fedavg
+  rounds: 50
+  local_steps: 250
+  min_clients: 3
+  execution: flower
+  save_every_n_rounds: 10
+  persist_client_state: true   # warm Adam per client; cold Adam every round
+                               # was a measured GR00T quality gap
+  client_state_keep: 2         # ~20GB per client per save; keep the last 2
+  output_dir: outputs/groot_cubestack_fl_3client