#!/bin/bash
# Run 5x5 grid eval for a given model type
# Usage: bash run_grid_eval.sh <model_type> <checkpoint> <out_dir> <gpu_id>

MODEL_TYPE=$1
CHECKPOINT=$2
OUT_DIR=$3
GPU_ID=$4

eval "$(conda shell.bash hook 2>/dev/null)"
conda activate motionlora

CENTER_DX=0.0509
CENTER_DY=-0.2063

# 5x5 grid
DX_VALUES=(-0.40 -0.3025 -0.205 -0.1075 -0.01)
DY_VALUES=(-0.30 -0.15 0.0 0.15 0.30)

cell=0
for ri in 0 1 2 3 4; do
    for ci in 0 1 2 3 4; do
        dx=$(echo "$CENTER_DX + ${DX_VALUES[$ri]}" | bc -l)
        dy=$(echo "$CENTER_DY + ${DY_VALUES[$ci]}" | bc -l)
        cell_dir="${OUT_DIR}/cell_${cell}"
        echo "=== Cell $cell: ri=$ri ci=$ci dx=$dx dy=$dy ==="

        CUDA_VISIBLE_DEVICES=$GPU_ID python /data/cameron/para_normalized_losses/libero/eval_multistage.py \
            --model_type $MODEL_TYPE \
            --checkpoint $CHECKPOINT \
            --benchmark libero_spatial --task_id 0 \
            --n_episodes 1 \
            --teleport --zero_rotation --clean_scene \
            --max_steps 600 \
            --shift_dx $dx --shift_dy $dy \
            --save_video \
            --out_dir $cell_dir

        cell=$((cell + 1))
    done
done

echo "=== All 25 cells done ==="