#!/bin/bash

# SBATCH params

conda activate radt

####### Train FetchReachObstacle #######

python3 train.py \
        --exp-name mg-radt-training-fetchreach \
        --instance_prefix run1 \
        --env reach_obstacle \
        --dataset_path 'INSERT/TRAINING/DATA/RELATIVE/PATH/HERE' \
        --n_head 4 \
        --n_layer 4 \
        --embed_dim 256 \
        --n_positions 1024 \
        --K 50 \
        --max_ep_len 50 \
        --batch_size 128 \
        --test_eval_interval 1000 \
        --max_iters 50000 \
        --scheduler 'cosinewarmuprestarts' \
        --T_0 10000 \
        --warmup_steps 5000 \
        --avoid_prompt \
        --max_avoid_prompt_len 10 \
        --adelta 2 \
        --alpha1 1.0 \
        --n_test_env 1 \
        --num_eval_episodes 60 \
        --bsa_box_size 0.08


####### Train MazeObstacle #######

# python3 train.py \
#         --exp-name mg-radt-training-maze \
#         --instance_prefix run1 \
#         --env pointmaze_obstacle \
#         --dataset_path 'INSERT/TRAINING/DATA/RELATIVE/PATH/HERE' \
#         --n_head 6 \
#         --n_layer 6 \
#         --embed_dim 384 \
#         --max_prompt_len 1 \
#         --K 300 \
#         --max_ep_len 300 \
#         --batch_size 32 \
#         --test_eval_interval 1000 \
#         --max_iters 50000 \
#         --avoid_prompt \
#         --max_avoid_prompt_len 10 \
#         --adelta 1 \
#         --alpha1 1 \
#         --alpha2 1.1 \
#         --buffer_size 0.2 \
#         --maze "U_MAZE" \
#         --n_test_env 1 \
#         --num_eval_episodes 100 \
#         --num_avoid 3



####### Train Cardiogenesis #######

# python3 train.py \
#         --exp-name mg-radt-training-cardio \
#         --instance_prefix run1 \
#         --env cardiogenesis \
#         --dataset_path 'INSERT/TRAINING/DATA/RELATIVE/PATH/HERE' \
#         --n_head 6 \
#         --n_layer 6 \
#         --embed_dim 384 \
#         --max_prompt_len 1 \
#         --K 30 \
#         --max_ep_len 30 \
#         --batch_size 128 \
#         --test_eval_interval 100 \
#         --save-interval 100 \
#         --max_iters 10000 \
#         --avoid_prompt \
#         --max_avoid_prompt_len 10 \
#         --adelta 1 \
#         --alpha1 1 \
#         --fixed_interval 10 \
#         --buffer_size 0.001
