#!/bin/bash
set -e

# ensure config name is set
if [ -z "$CONFIG_NAME" ]; then
    echo "CONFIG_NAME is not set"
    exit 1
fi


# compute norm stats
echo "Computing norm stats..."
python /opt/ml/code/openpi/scripts/compute_norm_stats.py --config-name $CONFIG_NAME

# train
echo "Training..."
python /opt/ml/code/openpi/scripts/train.py $CONFIG_NAME