From 441c7e4beafe1e50a5a2dda71d5563d2f5dd8541 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 15 Jul 2026 23:17:51 +0200 Subject: [PATCH] docs(robomme): support sample-matched multi-gpu ablations --- docs/source/robomme.mdx | 20 ++++++ .../robomme/smolvla_visual_memory_ablation.sh | 70 ++++++++++++++----- 2 files changed, 73 insertions(+), 17 deletions(-) diff --git a/docs/source/robomme.mdx b/docs/source/robomme.mdx index c93764b5f..9cf9e2c58 100644 --- a/docs/source/robomme.mdx +++ b/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 diff --git a/examples/robomme/smolvla_visual_memory_ablation.sh b/examples/robomme/smolvla_visual_memory_ablation.sh index 09d03b485..bb22d7fe6 100755 --- a/examples/robomme/smolvla_visual_memory_ablation.sh +++ b/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 \ - "${COMMON_TRAIN_ARGS[@]}" \ - --policy.use_visual_memory=false \ - --output_dir="${OUTPUT_ROOT}/baseline" \ - --job_name=robomme-smolvla-baseline - - uv run lerobot-train \ - "${COMMON_TRAIN_ARGS[@]}" \ - --policy.use_visual_memory=true \ - --policy.visual_memory_frames=6 \ - --policy.visual_memory_stride=10 \ - --policy.visual_memory_temporal_attention_every=4 \ - --output_dir="${OUTPUT_ROOT}/visual-memory" \ - --job_name=robomme-smolvla-visual-memory + for variant in "${VARIANTS[@]}"; do + MEMORY_ARGS=(--policy.use_visual_memory=false) + if [[ "${variant}" == "visual-memory" ]]; then + MEMORY_ARGS=( + --policy.use_visual_memory=true + --policy.visual_memory_frames=6 + --policy.visual_memory_stride=10 + --policy.visual_memory_temporal_attention_every=4 + ) + fi + "${TRAIN_COMMAND[@]}" \ + "${COMMON_TRAIN_ARGS[@]}" \ + "${MEMORY_ARGS[@]}" \ + --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