#!/bin/bash
set -x
cd /workspace/rl4phyx

MODEL_PATH="/workspace/rl4phyx/models/Qwen2.5-VL-3B-Instruct"
BASE="/workspace/rl4phyx/RL4Phyx/ZeroSearch/One-Shot-RLVR"
DATA_FILE="${BASE}/data/train/physics_vlm/mechanics/mechanics_1_rl_numerical.parquet"
VAL_FILE="/workspace/rl4phyx/RL4Phyx/oneshot/validation_data/metaphyx_oe_1533.parquet"
IMAGE_DIR="${BASE}/data/train/physics_vlm/mechanics"

export VLLM_ATTENTION_BACKEND=XFORMERS
export VLLM_USE_TRITON_FLASH_ATTN=0
export CUDA_VISIBLE_DEVICES=0,1,2,3

python3 -m verl.trainer.main_ppo \
 algorithm.adv_estimator=grpo \
 +data_domain=physics \
 data.train_files=${DATA_FILE} \
 data.val_files=${VAL_FILE} \
 data.train_batch_size=128 \
 data.val_batch_size=1533 \
 data.max_prompt_length=2048 \
 data.max_response_length=3072 \
 +data.is_multimodal=True \
 +data.vlm_model=${MODEL_PATH} \
 +data.image_dir=${IMAGE_DIR} \
 reward_model.reward_manager='naive' \
 actor_rollout_ref.model.path=${MODEL_PATH} \
 actor_rollout_ref.actor.optim.lr=1e-6 \
 actor_rollout_ref.model.use_remove_padding=False \
 actor_rollout_ref.actor.ppo_mini_batch_size=128 \
 actor_rollout_ref.actor.use_dynamic_bsz=True \
 actor_rollout_ref.actor.ppo_max_token_len_per_gpu=24000 \
 actor_rollout_ref.actor.use_kl_loss=True \
 actor_rollout_ref.actor.kl_loss_coef=0.001 \
 actor_rollout_ref.actor.kl_loss_type=low_var_kl \
 actor_rollout_ref.model.enable_gradient_checkpointing=True \
 actor_rollout_ref.actor.fsdp_config.param_offload=False \
 +actor_rollout_ref.actor.fsdp_config.grad_offload=False \
 actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
 actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
 actor_rollout_ref.rollout.name=vllm \
 actor_rollout_ref.rollout.temperature=0.7 \
 +actor_rollout_ref.rollout.val_temperature=0.7 \
 actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \
 actor_rollout_ref.rollout.n=8 \
 +actor_rollout_ref.rollout.n_val=1 \
 actor_rollout_ref.ref.fsdp_config.param_offload=True \
 algorithm.kl_ctrl.kl_coef=0.001 \
 trainer.critic_warmup=0 \
 trainer.logger=['console'] \
 trainer.project_name=physics_rlvr_dryrun \
 trainer.experiment_name=dryrun_baseline \
 trainer.checkpoints_dir=/workspace/rl4phyx/checkpoints/dryrun \
 +trainer.val_before_train=False \
 trainer.n_gpus_per_node=4 \
 trainer.nnodes=1 \
 trainer.save_freq=100 \
 trainer.test_freq=100 \
 trainer.default_hdfs_dir=null \
 trainer.total_epochs=5 2>&1 | tee /workspace/rl4phyx/logs/dryrun_baseline.log