#!/usr/bin/sh

DIR="experiment_runs/diffmorl"
# expert
for ENV in MO-Ant-v2 MO-HalfCheetah-v2 MO-Hopper-v2 MO-Hopper-v3 MO-Swimmer-v2 MO-Walker2d-v2
do
    if [ "$ENV" = "MO-Ant-v2" ]; then
        STEP=100000
        MNUM=8
        MSTEP=100000
    elif [ "$ENV" = "MO-HalfCheetah-v2" ]; then
        STEP=400000
        MNUM=6
        MSTEP=400000
    elif [ "$ENV" = "MO-Hopper-v2" ]; then
        STEP=400000
        MNUM=6
        MSTEP=50000
    elif [ "$ENV" = "MO-Hopper-v3" ]; then
        STEP=200000
        MNUM=5
        MSTEP=100000
    elif [ "$ENV" = "MO-Swimmer-v2" ]; then
        STEP=200000
        MNUM=5
        MSTEP=100000
    elif [ "$ENV" = "MO-Walker2d-v2" ]; then
        STEP=400000
        MNUM=6
        MSTEP=150000
    fi

    for QUALITY in expert
    do
        for SEED in 1 2 3
        do  
        {
            CUDA_VISIBLE_DEVICES=0 python experiment.py --dir $DIR --env $ENV --seed $SEED --dataset $QUALITY'_'uniform --model_type mod --mod_type bc --num_steps_per_iter $STEP --max_iters 1 --use_p_bar True --K 8 --infer_N 7 --n_diffusion_steps 8 --returns_condition True --mixup True --mixup_num $MNUM --mixup_step $MSTEP &
            CUDA_VISIBLE_DEVICES=1 python experiment.py --dir $DIR --env $ENV --seed $SEED --dataset $QUALITY'_'wide --model_type mod --mod_type bc --num_steps_per_iter $STEP --max_iters 1 --use_p_bar True --K 8 --infer_N 7 --n_diffusion_steps 8 --returns_condition True --mixup True --mixup_num $MNUM --mixup_step $MSTEP &
            CUDA_VISIBLE_DEVICES=0 python experiment.py --dir $DIR --env $ENV --seed $SEED --dataset $QUALITY'_'narrow --model_type mod --mod_type bc --num_steps_per_iter $STEP --max_iters 1 --use_p_bar True --K 8 --infer_N 7 --n_diffusion_steps 8 --returns_condition True --mixup True --mixup_num $MNUM --mixup_step $MSTEP &
            CUDA_VISIBLE_DEVICES=1 python experiment.py --dir $DIR --env $ENV --seed $SEED --dataset $QUALITY'_'custom --model_type mod --mod_type bc --num_steps_per_iter $STEP --max_iters 1 --use_p_bar True --K 8 --infer_N 7 --n_diffusion_steps 8 --returns_condition True --mixup True --mixup_num $MNUM --mixup_step $MSTEP &
            wait
        }
        done
    done
done

# amateur 
for ENV in MO-Ant-v2 MO-HalfCheetah-v2 MO-Hopper-v2 MO-Hopper-v3 MO-Swimmer-v2 MO-Walker2d-v2
do
    CFG_W=0.1
    if [ "$ENV" = "MO-Ant-v2" ]; then
        STEP=100000
        MNUM=6
        MSTEP=100000
    elif [ "$ENV" = "MO-HalfCheetah-v2" ]; then
        STEP=200000
        MNUM=6
        MSTEP=200000
        CFG_W=1
    elif [ "$ENV" = "MO-Hopper-v2" ]; then
        STEP=300000
        MNUM=6
        MSTEP=200000
    elif [ "$ENV" = "MO-Hopper-v3" ]; then
        STEP=100000
        MNUM=5
        MSTEP=100000
    elif [ "$ENV" = "MO-Swimmer-v2" ]; then
        STEP=50000
        MNUM=5
        MSTEP=50000
    elif [ "$ENV" = "MO-Walker2d-v2" ]; then
        STEP=100000
        MNUM=6
        MSTEP=100000
        CFG_W=1
    fi

    for QUALITY in amateur
    do
        for SEED in 1 2 3
        do  
        {
            CUDA_VISIBLE_DEVICES=0 python experiment.py --dir $DIR --env $ENV --seed $SEED --dataset $QUALITY'_'uniform --model_type mod --mod_type bc --num_steps_per_iter $STEP --max_iters 1 --use_p_bar True --K 8 --infer_N 7 --n_diffusion_steps 8 --returns_condition True --mixup True --mixup_num $MNUM --mixup_step $MSTEP --v_cfg_W $CFG_W &
            CUDA_VISIBLE_DEVICES=1 python experiment.py --dir $DIR --env $ENV --seed $SEED --dataset $QUALITY'_'wide --model_type mod --mod_type bc --num_steps_per_iter $STEP --max_iters 1 --use_p_bar True --K 8 --infer_N 7 --n_diffusion_steps 8 --returns_condition True --mixup True --mixup_num $MNUM --mixup_step $MSTEP --v_cfg_W $CFG_W &
            CUDA_VISIBLE_DEVICES=0 python experiment.py --dir $DIR --env $ENV --seed $SEED --dataset $QUALITY'_'narrow --model_type mod --mod_type bc --num_steps_per_iter $STEP --max_iters 1 --use_p_bar True --K 8 --infer_N 7 --n_diffusion_steps 8 --returns_condition True --mixup True --mixup_num $MNUM --mixup_step $MSTEP --v_cfg_W $CFG_W &
            CUDA_VISIBLE_DEVICES=1 python experiment.py --dir $DIR --env $ENV --seed $SEED --dataset $QUALITY'_'custom --model_type mod --mod_type bc --num_steps_per_iter $STEP --max_iters 1 --use_p_bar True --K 8 --infer_N 7 --n_diffusion_steps 8 --returns_condition True --mixup True --mixup_num $MNUM --mixup_step $MSTEP --v_cfg_W $CFG_W &
            wait
        }
        done
    done
done