#!/bin/bash
# The ACT workshop's GPU runs, exactly as made for the page (RTX 4060 Ti; times measured there).
# Run from anywhere; everything is written to ~/lerobot-runs/w4b (about 3 GB, outside the repository).
#   train: 10,000 steps, batch 8, about 23 min per seed; evaluate: 200 starts, about 9 min per policy.
set -e
mkdir -p ~/lerobot-runs/w4b && cd ~/lerobot-runs/w4b
export MUJOCO_GL=egl
L=~/.virtualenvs/lerobot/bin

train () {  # $1 seed, $2 output folder
  $L/lerobot-train --policy.type=act --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human --dataset.video_backend=pyav \
    --steps=10000 --batch_size=8 --log_freq=250 --save_freq=2500 --policy.push_to_hub=false --policy.device=cuda \
    --output_dir=$2 --num_workers=8 --seed=$1 > $2.log 2>&1
}
evaluate () {  # $1 policy folder, $2 output name, then any policy overrides
  policy=$1; out=$2; shift 2
  $L/lerobot-eval --policy.path=$policy "$@" --env.type=aloha --env.task=AlohaTransferCube-v0 --eval.n_episodes=200 \
    --eval.batch_size=10 --eval.use_async_envs=false --policy.device=cuda --seed=1000 --output_dir=$out > $out.log 2>&1
}

# The published policy (100,000 steps). Saved by an older LeRobot, so its normalisation is moved into processor files once.
$L/python -m lerobot.processor.migrate_policy_normalization --pretrained-path lerobot/act_aloha_sim_transfer_cube_human \
  --output-dir published_migrated
evaluate published_migrated eval200_published

# Our two training runs, each tested at four checkpoints on the same 200 starts.
train 1000 train_act
train 2000 train_act_seed2
for c in 002500 005000 007500 010000; do
  evaluate train_act/checkpoints/$c/pretrained_model eval200_$c
  evaluate train_act_seed2/checkpoints/$c/pretrained_model eval200_seed2_$c
done

# One thing varied: how often the published policy looks again (no retraining).
evaluate published_migrated eval200_pub_n50 --policy.n_action_steps=50
evaluate published_migrated eval200_pub_n25 --policy.n_action_steps=25
evaluate published_migrated eval200_pub_ens --policy.n_action_steps=1 --policy.temporal_ensemble_coeff=0.01