docs(robomme): support sample-matched multi-gpu ablations

docs/source/robomme.mdx

@@ -145,6 +145,26 @@ execution boundary while preserving earlier frames for temporal observation delt
 training split this keeps 476,857 execution targets out of 768,897 total frames. One execution-target
 epoch therefore takes `ceil(476857 / effective_batch_size)` optimizer steps.
 
+For a sample-matched SmolVLA visual-memory ablation, use
+`examples/robomme/smolvla_visual_memory_ablation.sh`. `TARGET_SAMPLES` counts examples across all
+GPUs, and `NUM_PROCESSES` is included when the script converts that target into optimizer steps. For
+example, the following reproduces 5.12 million example exposures (the exposure of RoboMME's
+80,000-step, global-batch-64 memory-policy recipe) with four GPUs and a global batch of 48:
+
+```bash
+TARGET_SAMPLES=5120000 \
+NUM_PROCESSES=4 \
+BATCH_SIZE=12 \
+VARIANT=visual-memory \
+RUN_EVAL=false \
+bash examples/robomme/smolvla_visual_memory_ablation.sh
+```
+
+This becomes 106,667 optimizer steps, or about 10.74 execution-target epochs. Run the baseline with
+the same `TARGET_SAMPLES`, effective batch size, seed, and scheduler settings to isolate the visual
+memory change. RoboMME's released baseline uses a different global batch (128), so its nominal
+80,000-step recipe is not sample-matched to its global-batch-64 memory recipe.
+
 At evaluation time the environment wrapper has already converted observations to canonical keys. Only the two camera suffixes need to be aligned with the checkpoint:
 
 ```bash

examples/robomme/smolvla_visual_memory_ablation.sh

@@ -5,18 +5,48 @@
 set -euo pipefail
 
 BATCH_SIZE="${BATCH_SIZE:-4}"
+NUM_PROCESSES="${NUM_PROCESSES:-1}"
 # The published training split has 476,857 execution frames. By default, train on one
 # execution-frame epoch; set STEPS explicitly to use a different optimizer-step budget.
 TARGET_SAMPLES="${TARGET_SAMPLES:-476857}"
-STEPS="${STEPS:-$(((TARGET_SAMPLES + BATCH_SIZE - 1) / BATCH_SIZE))}"
+EFFECTIVE_BATCH_SIZE=$((BATCH_SIZE * NUM_PROCESSES))
+STEPS="${STEPS:-$(((TARGET_SAMPLES + EFFECTIVE_BATCH_SIZE - 1) / EFFECTIVE_BATCH_SIZE))}"
+SCHEDULER_WARMUP_STEPS="${SCHEDULER_WARMUP_STEPS:-$(((STEPS + 29) / 30))}"
+SCHEDULER_DECAY_STEPS="${SCHEDULER_DECAY_STEPS:-${STEPS}}"
 SEED="${SEED:-1000}"
 OUTPUT_ROOT="${OUTPUT_ROOT:-outputs/robomme-smolvla-mem-ablation}"
 WANDB_ENABLE="${WANDB_ENABLE:-false}"
 RUN_TRAIN="${RUN_TRAIN:-true}"
 RUN_EVAL="${RUN_EVAL:-true}"
+VARIANT="${VARIANT:-both}"
 TASKS="BinFill,PickXtimes,SwingXtimes,StopCube,VideoUnmask,VideoUnmaskSwap,ButtonUnmask,ButtonUnmaskSwap,PickHighlight,VideoRepick,VideoPlaceButton,VideoPlaceOrder,MoveCube,InsertPeg,PatternLock,RouteStick"
 TASK_IDS="[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49]"
 
+case "${VARIANT}" in
+  baseline) VARIANTS=(baseline) ;;
+  visual-memory) VARIANTS=(visual-memory) ;;
+  both) VARIANTS=(baseline visual-memory) ;;
+  *) echo "VARIANT must be baseline, visual-memory, or both" >&2; exit 2 ;;
+esac
+
+TRAIN_COMMAND=()
+if [[ "${RUN_TRAIN}" == "true" ]]; then
+  if ((NUM_PROCESSES > 1)); then
+    TRAIN_ENTRYPOINT="$(uv run which lerobot-train)"
+    TRAIN_COMMAND=(
+      uv run accelerate launch
+      --multi_gpu
+      --num_processes="${NUM_PROCESSES}"
+      --mixed_precision=bf16
+      "${TRAIN_ENTRYPOINT}"
+    )
+  else
+    TRAIN_COMMAND=(uv run lerobot-train)
+  fi
+fi
+
+echo "Training ${VARIANT}: ${TARGET_SAMPLES} target samples, global batch ${EFFECTIVE_BATCH_SIZE}, ${STEPS} optimizer steps"
+
 COMMON_TRAIN_ARGS=(
   --policy.path=lerobot/smolvla_base
   --policy.device=cuda
@@ -29,6 +59,8 @@ COMMON_TRAIN_ARGS=(
   '--rename_map={"image":"observation.images.camera1","wrist_image":"observation.images.camera2","state":"observation.state","actions":"action"}'
   --batch_size="${BATCH_SIZE}"
   --steps="${STEPS}"
+  --policy.scheduler_warmup_steps="${SCHEDULER_WARMUP_STEPS}"
+  --policy.scheduler_decay_steps="${SCHEDULER_DECAY_STEPS}"
   --seed="${SEED}"
   --env_eval_freq=0
   --save_freq=5000
@@ -36,24 +68,26 @@ COMMON_TRAIN_ARGS=(
 )
 
 if [[ "${RUN_TRAIN}" == "true" ]]; then
-  uv run lerobot-train \
+  for variant in "${VARIANTS[@]}"; do
-    "${COMMON_TRAIN_ARGS[@]}" \
+    MEMORY_ARGS=(--policy.use_visual_memory=false)
-    --policy.use_visual_memory=false \
+    if [[ "${variant}" == "visual-memory" ]]; then
-    --output_dir="${OUTPUT_ROOT}/baseline" \
+      MEMORY_ARGS=(
-    --job_name=robomme-smolvla-baseline
+        --policy.use_visual_memory=true
-
+        --policy.visual_memory_frames=6
-  uv run lerobot-train \
+        --policy.visual_memory_stride=10
-    "${COMMON_TRAIN_ARGS[@]}" \
+        --policy.visual_memory_temporal_attention_every=4
-    --policy.use_visual_memory=true \
+      )
-    --policy.visual_memory_frames=6 \
+    fi
-    --policy.visual_memory_stride=10 \
+    "${TRAIN_COMMAND[@]}" \
-    --policy.visual_memory_temporal_attention_every=4 \
+      "${COMMON_TRAIN_ARGS[@]}" \
-    --output_dir="${OUTPUT_ROOT}/visual-memory" \
+      "${MEMORY_ARGS[@]}" \
-    --job_name=robomme-smolvla-visual-memory
+      --output_dir="${OUTPUT_ROOT}/${variant}" \
+      --job_name="robomme-smolvla-${variant}"
+  done
 fi
 
 if [[ "${RUN_EVAL}" == "true" ]]; then
-  for variant in baseline visual-memory; do
+  for variant in "${VARIANTS[@]}"; do
     uv run lerobot-eval \
       --policy.path="${OUTPUT_ROOT}/${variant}/checkpoints/last/pretrained_model" \
       --env.type=robomme \
@@ -68,7 +102,8 @@ if [[ "${RUN_EVAL}" == "true" ]]; then
   done
 fi
 
-uv run python - "${OUTPUT_ROOT}" <<'PY'
+if [[ "${RUN_EVAL}" == "true" ]]; then
+  uv run python - "${OUTPUT_ROOT}" <<'PY'
 import json
 import pathlib
 import sys
@@ -91,3 +126,4 @@ for variant in ("baseline", "visual-memory"):
             f"avg_reward={metrics['avg_sum_reward']:.4f}"
         )
 PY
+fi