mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-27 11:46:04 +00:00
Pepijn
94 lines
3.4 KiB
Bash
Executable File
94 lines
3.4 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# Matched SmolVLA ablation for the 10 fps lerobot/robomme dataset.
|
|
# Run from the LeRobot repository root on a CUDA machine.
|
|
|
|
set -euo pipefail
|
|
|
|
BATCH_SIZE="${BATCH_SIZE:-4}"
|
|
# 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))}"
|
|
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}"
|
|
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]"
|
|
|
|
COMMON_TRAIN_ARGS=(
|
|
--policy.path=lerobot/smolvla_base
|
|
--policy.device=cuda
|
|
--policy.push_to_hub=false
|
|
--policy.empty_cameras=1
|
|
--policy.freeze_vision_encoder=false
|
|
--policy.train_expert_only=false
|
|
--dataset.repo_id=lerobot/robomme
|
|
--dataset.training_target_start_feature=exec_start_idx
|
|
'--rename_map={"image":"observation.images.camera1","wrist_image":"observation.images.camera2","state":"observation.state","actions":"action"}'
|
|
--batch_size="${BATCH_SIZE}"
|
|
--steps="${STEPS}"
|
|
--seed="${SEED}"
|
|
--env_eval_freq=0
|
|
--save_freq=5000
|
|
--wandb.enable="${WANDB_ENABLE}"
|
|
)
|
|
|
|
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
|
|
fi
|
|
|
|
if [[ "${RUN_EVAL}" == "true" ]]; then
|
|
for variant in baseline visual-memory; do
|
|
uv run lerobot-eval \
|
|
--policy.path="${OUTPUT_ROOT}/${variant}/checkpoints/last/pretrained_model" \
|
|
--env.type=robomme \
|
|
--env.task="${TASKS}" \
|
|
--env.dataset_split=test \
|
|
--env.task_ids="${TASK_IDS}" \
|
|
'--rename_map={"observation.images.image":"observation.images.camera1","observation.images.wrist_image":"observation.images.camera2"}' \
|
|
--eval.batch_size=1 \
|
|
--eval.n_episodes=1 \
|
|
--seed="${SEED}" \
|
|
--output_dir="${OUTPUT_ROOT}/eval-${variant}"
|
|
done
|
|
fi
|
|
|
|
uv run python - "${OUTPUT_ROOT}" <<'PY'
|
|
import json
|
|
import pathlib
|
|
import sys
|
|
|
|
root = pathlib.Path(sys.argv[1])
|
|
for variant in ("baseline", "visual-memory"):
|
|
result_path = root / f"eval-{variant}" / "eval_info.json"
|
|
if not result_path.exists():
|
|
continue
|
|
with result_path.open() as handle:
|
|
info = json.load(handle)
|
|
overall = info["overall"]
|
|
print(
|
|
f"{variant}: success={overall['pc_success']:.2f}% "
|
|
f"avg_reward={overall['avg_sum_reward']:.4f}"
|
|
)
|
|
for task, metrics in info["per_group"].items():
|
|
print(
|
|
f" {task}: success={metrics['pc_success']:.2f}% "
|
|
f"avg_reward={metrics['avg_sum_reward']:.4f}"
|
|
)
|
|
PY
|