################ We collect data using TS, and optimal 
################ action is the next action (lookahead)

# # # Collect data using unif
# CUDA_VISIBLE_DEVICES=1 python3 collect_data_lookahead.py --env linear_bandit_train_lookahead_pred_reward \
#     --envs 100001 \
#     --H 202 \
#     --dim 5 \
#     --var 0.3 \
#     --cov 0.0 \
#     --lin_d 2 \
#     --envs_eval 200 \
#     --data_type unif_pred_reward \
#     --pred_reward_type linear

# # # Train
# CUDA_VISIBLE_DEVICES=1 python3 train_new_wt.py --env linear_bandit_train_lookahead_pred_reward \
#     --envs 100001 \
#     --H 202 \
#     --dim 5 \
#     --lin_d 2 \
#     --var 0.3 \
#     --cov 0.0 \
#     --lr 0.00015 \
#     --layer 4 \
#     --head 4 \
#     --num_epochs 1000 \
#     --batch_size 1024 \
#     --seed 1 \
#     --pred_reward_type linear

# Evaluate, choose an appropriate epoch
CUDA_VISIBLE_DEVICES=1 python3 eval_new.py --env linear_bandit_train_lookahead_pred_reward \
    --envs 100001 \
    --H 202 \
    --dim 5 \
    --lin_d 2 \
    --var 0.3 \
    --cov 0.0 \
    --lr 0.00015 \
    --layer 4 \
    --head 4 \
    --epoch 1000 \
    --n_eval 200 \
    --seed 1 \
    --pred_reward_type linear
