# Run from repo root
cd "$(dirname "$0")/.."
REPO_ROOT="$(pwd)"
export PYTHONPATH="${REPO_ROOT}:${PYTHONPATH}"

# DROID TFDS: use official GCS bucket (requires gcloud auth) or a local data_dir where
# you have built/downloaded the DROID TFDS dataset (see https://droid-dataset.github.io/)
DATA_DIR="${DROID_DATA_DIR:-gs://gresearch/robotics}"
BUFFER_PATH="${REPO_ROOT}/data/buffer_droid.zarr"
mkdir -p "$(dirname "$BUFFER_PATH")"
if [ ! -d "$BUFFER_PATH" ]; then
  # Cache dataset (minimal: 50 episodes for a quick run)
  echo "Caching dataset from $DATA_DIR..."
  python datasets/droid/convert_dataset_zarr.py --data_dir "$DATA_DIR" --buffer_path "$BUFFER_PATH" --num_episodes 50 --num_workers 4
fi

# UWM training
python experiments/uwm/train.py dataset=droid exp_id=minimal dataset.buffer_path=$BUFFER_PATH

# DP
# python experiments/dp/train.py dataset=droid exp_id=benchmark dataset.buffer_path=$BUFFER_PATH

# GR1
# python experiments/gr1/train.py dataset=droid exp_id=benchmark dataset.buffer_path=$BUFFER_PATH

# PAD
# python experiments/pad/train.py dataset=droid exp_id=benchmark dataset.buffer_path=$BUFFER_PATH