Batch 16, no gradient checkpointing, 340k steps (1 epoch)

train_cloud.sh

@@ -35,15 +35,15 @@ python3.12 -m accelerate.commands.launch $ACCEL_FLAGS \
 --policy.path=lerobot/pi05_base \
 --policy.train_expert_only=true \
 --policy.dtype=bfloat16 \
---policy.gradient_checkpointing=true \
+--policy.gradient_checkpointing=false \
 --policy.push_to_hub=true \
 --policy.repo_id=StrongRoboticsLab/pi05-so100-diverse \
 --policy.normalization_mapping='{"VISUAL": "IDENTITY", "STATE": "MEAN_STD", "ACTION": "MEAN_STD"}' \
 --policy.scheduler_warmup_steps=1000 \
---policy.scheduler_decay_steps=170000 \
+--policy.scheduler_decay_steps=340000 \
 --rename_map='{"observation.images.image": "observation.images.base_0_rgb", "observation.images.image2": "observation.images.left_wrist_0_rgb"}' \
---batch_size=32 \
---steps=170000 \
+--batch_size=16 \
+--steps=340000 \
 --save_freq=500 \
 --log_freq=50 \
 --num_workers=4 \