################ We collect data using TS, and optimal 
################ action is the next action (lookahead)
################ First create dim-1 arms same, then add 1 new arm randomly to each new env

# # Collect data using TS
CUDA_VISIBLE_DEVICES=1 python3 collect_data_lookahead.py --env linear_bandit_train_original_emp_opt \
    --envs 200002 \
    --H 51 \
    --dim 20 \
    --var 0.3 \
    --cov 0.0 \
    --lin_d 5 \
    --envs_eval 200 \
    --data_type TS_original \
    --pred_reward_type new_arms_linear \
    --new_arms 1

# # Train
CUDA_VISIBLE_DEVICES=1 python3 train_new_wt.py --env linear_bandit_train_original_emp_opt \
    --envs 200002 \
    --H 51 \
    --dim 20 \
    --lin_d 5 \
    --var 0.3 \
    --cov 0.0 \
    --lr 0.00015 \
    --layer 4 \
    --head 4 \
    --num_epochs 1000 \
    --batch_size 1024 \
    --seed 1 \
    --pred_reward_type new_arms_linear

# # Evaluate, choose an appropriate epoch
CUDA_VISIBLE_DEVICES=1 python3 eval_new.py --env linear_bandit_train_original_emp_opt \
    --envs 200002 \
    --H 51 \
    --dim 20 \
    --lin_d 5 \
    --var 0.3 \
    --cov 0.0 \
    --lr 0.00015 \
    --layer 4 \
    --head 4 \
    --epoch 950 \
    --n_eval 200 \
    --seed 1 \
    --pred_reward_type new_arms_linear
