Batch 32, 170k steps (1 epoch), no gradient checkpointing

train_cloud.sh

@@ -40,10 +40,10 @@ python3.12 -m accelerate.commands.launch $ACCEL_FLAGS \
 --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=85000 \
+--policy.scheduler_decay_steps=170000 \
 --rename_map='{"observation.images.image": "observation.images.base_0_rgb", "observation.images.image2": "observation.images.left_wrist_0_rgb"}' \
---batch_size=64 \
---steps=85000 \
+--batch_size=32 \
+--steps=170000 \
 --save_freq=500 \
 --log_freq=50 \
 --num_workers=4 \