PROJECT_ROOT="./your/path/to/DiagRL"
export PYTHONUNBUFFERED=1
export HYDRA_FULL_ERROR=1
export CUDA_VISIBLE_DEVICES=0,1,2,3
# 将src目录添加到PYTHONPATH
export PYTHONPATH="${PROJECT_ROOT}/src:${PYTHONPATH}"

NGPUS=4
MODEL_PARALLEL=4
MATCH_TOP_N=20
SEARCH_MAX_N=10
PROJECT_NAME=DiagRL
EXPERIMENT_NAME=DiagRLMIMIC
TRAIN_FILES=./your/path/to/DiagRL/data/train_data_MIMICC.parquet
TEST_FILES=./your/path/to/DiagRL/data/val_data_MIMICC.parquet
TRAIN_BATCH_SIZE=64
PROMPT_KEY=input
MAX_PROMPT_LENGTH=1024
MAX_RESPONSE_LENGTH=6144
PROMPT_TEMPLATE_NAME=rare_search_template_sys
APPLY_CHAT=True
ACTOR_MODEL_PATH=Qwen/Qwen2.5-7B-Instruct
PPO_MINI_BATCH_SIZE=16
PPO_MICRO_BATCH_SIZE=4
REWARD_MANAGER=rare_search
TOTAL_TRAINING_STEPS=1000

python3 -m verl.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    data.train_files="$TRAIN_FILES" \
    data.val_files="$TEST_FILES" \
    data.prompt_key=${PROMPT_KEY} \
    data.train_batch_size=${TRAIN_BATCH_SIZE} \
    data.val_batch_size=${TRAIN_BATCH_SIZE} \
    +data.shuffle_train_dataloader=True \
    data.max_prompt_length=${MAX_PROMPT_LENGTH} \
    data.max_response_length=${MAX_RESPONSE_LENGTH} \
    data.apply_chat=${APPLY_CHAT} \
    data.prompt_template_name=${PROMPT_TEMPLATE_NAME} \
    actor_rollout_ref.model.path=${ACTOR_MODEL_PATH} \
    actor_rollout_ref.model.use_remove_padding=True \
    actor_rollout_ref.model.enable_gradient_checkpointing=true \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.285 \
    actor_rollout_ref.actor.use_kl_loss=true \
    actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE} \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${PPO_MICRO_BATCH_SIZE} \
    +actor_rollout_ref.rollout.match_source=./your/path/to/DiagRL/src/match/match_source_MIMICC_filtered.pt \
    +actor_rollout_ref.rollout.match_top_n=${MATCH_TOP_N} \
    +actor_rollout_ref.rollout.search_source=./your/path/to/DiagRL/src/search/common_disease_phenotype.json \
    +actor_rollout_ref.rollout.search_max_n=${SEARCH_MAX_N} \
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
    actor_rollout_ref.rollout.tensor_model_parallel_size=${MODEL_PARALLEL} \
    actor_rollout_ref.rollout.name=vllm_with_search \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
    actor_rollout_ref.ref.fsdp_config.param_offload=True \
    actor_rollout_ref.actor.kl_loss_coef=0.001 \
    actor_rollout_ref.actor.kl_loss_type=low_var_kl \
    reward_model.reward_manager=${REWARD_MANAGER} \
    algorithm.kl_ctrl.kl_coef=0.001 \
    trainer.logger=['console'] \
    trainer.project_name=${PROJECT_NAME} \
    trainer.experiment_name=${EXPERIMENT_NAME} \
    trainer.val_before_train=False \
    trainer.default_hdfs_dir=null \
    trainer.n_gpus_per_node=${NGPUS} \
    trainer.nnodes=1 \
    trainer.save_freq=50 \
    trainer.test_freq=50 \
    trainer.total_training_steps=${TOTAL_TRAINING_STEPS} \
    trainer.rollout_save_path=./your/path/to/DiagRL/scripts/train/checkpoints/${PROJECT_NAME}/${EXPERIMENT_NAME}/outputs \
    trainer.total_epochs=15 2>&1 | tee verl_demo_${EXPERIMENT_NAME}.log