#!/bin/bash
# One-Shot RLVR: Physics (VLM) - Corrected for absolute paths
# Based on training_physics_mechanics_oe.sh (canonical VLM adaptation)
# Usage: bash run_rlvr.sh <MODEL_PATH> <EXPERIMENT_NAME>

set -x

# ==================== Accept CLI args ====================
MODEL_PATH="${1:-/workspace/rl4phyx/models/Qwen2.5-VL-3B-Instruct}"
EXPERIMENT_NAME="${2:-rlvr_baseline}"

# ==================== Configuration ====================
WANDB_PROJECT=${WANDB_PROJECT:-"physics_one_shot_rlvr"}
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"

# GPU
N_GPUS=4
GPUS="0,1,2,3"

# Training Hyperparameters (aligned with original One-Shot RLVR)
BATCH_SIZE="128"
LEARNING_RATE="1e-6"
ROLLOUT_N="8"
TEMPERATURE="0.7"
TENSOR_PARALLEL="2"

# Checkpoints
CHECKPOINT_DIR="/workspace/rl4phyx/checkpoints/${EXPERIMENT_NAME}"

echo "=========================================="
echo "One-Shot RLVR: Physics (VLM)"
echo "Model: ${MODEL_PATH}"
echo "Experiment: ${EXPERIMENT_NAME}"
echo "Data: ${DATA_FILE}"
echo "Checkpoint: ${CHECKPOINT_DIR}"
echo "=========================================="

mkdir -p $CHECKPOINT_DIR
mkdir -p /workspace/rl4phyx/logs

# IMPORTANT: Run from /workspace/rl4phyx so Python imports pip veRL v0.7
# NOT from One-Shot-RLVR dir (which has local verl/ v0.2)
cd /workspace/rl4phyx

# vLLM backend
export VLLM_ATTENTION_BACKEND=XFORMERS
export VLLM_USE_TRITON_FLASH_ATTN=0
export WANDB_API_KEY=${WANDB_API_KEY:-""}

# ==================== Training ====================
CUDA_VISIBLE_DEVICES=${GPUS} 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=${BATCH_SIZE} \
 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=${LEARNING_RATE} \
 actor_rollout_ref.model.use_remove_padding=False \
 actor_rollout_ref.actor.ppo_mini_batch_size=${BATCH_SIZE} \
 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=${TENSOR_PARALLEL} \
 actor_rollout_ref.rollout.name=vllm \
 actor_rollout_ref.rollout.temperature=${TEMPERATURE} \
 +actor_rollout_ref.rollout.val_temperature=${TEMPERATURE} \
 actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \
 actor_rollout_ref.rollout.n=${ROLLOUT_N} \
 +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','wandb'] \
 trainer.project_name=${WANDB_PROJECT} \
 trainer.experiment_name=${EXPERIMENT_NAME} \
 trainer.checkpoints_dir=$CHECKPOINT_DIR \
 +trainer.val_before_train=True \
 trainer.n_gpus_per_node=${N_GPUS} \
 trainer.nnodes=1 \
 trainer.save_freq=20 \
 trainer.test_freq=20 \
 trainer.default_hdfs_dir=null \
 trainer.total_epochs=2000 2>&1 | tee /workspace/rl4phyx/logs/${EXPERIMENT_NAME}.log