Commit 542f4698 authored by jameskrw's avatar jameskrw
Browse files

Merge branch 'debug' of github.com:JamesKrW/vagen into debug

parents 23395410 94ca316e
Loading
Loading
Loading
Loading
+2 −0
Original line number Diff line number Diff line
from vagen.env.register import REGISTERED_ENVS, register

from vagen.env.sokoban.env import SokobanInterface
 No newline at end of file
+13 −8
Original line number Diff line number Diff line
@@ -16,10 +16,11 @@ Rules:
2. Avoid walls (#).

Actions you can take: Up, Down, Left, Right. You can take up to {max_action_per_step} action(s) at a time.
Up: move up to the cell above (to the above row).
Down: move down to the cell below (to the below row).
Left: move left to the cell to the left (to the left column).
Right: move right to the cell to the right (to the right column).
- Up: move up to the cell above
- Down: move down to the cell below
- Left: move left to the cell to the left
- Right: move right to the cell to the right
If there is a box on the cell you want to move to, you will push the box one cell in the same direction.

Rewards:
Box on target: +1.0
@@ -27,7 +28,9 @@ All boxes placed: +10.0
Format correct: +0.5

Please think step by step and provide the actions you want to take.
Your reponse should be in the format of <think>...</think><answer>...</answer>
You should wrap your thought between `<think>` and `</think>` tags, and wrap your answer between `<answer>` and `</answer>` tags.
Your response should STRICTLY follow the format:
<think>...</think><answer>...</answer>
"""
# E.g. <think> There's a box on the upper right of me, the target is on the upper side of the box, I need to push the box it upward. </think><answer> Right,Up,Up </answer>
# Let's try to use a format reward and answer reward
@@ -40,14 +43,16 @@ init_observation_template = """
[Initial Observation]:
{observation}
Decide your next action(s).
Your reponse should be in the format of <think>...</think><answer>...</answer>
Your response should STRICTLY follow the format:
<think>...</think><answer>...</answer>
"""

action_template = """After your answer, the extracted valid action is {valid_action}.\
action_template = """Valid action extracted from your response is {valid_action}.\
After that, the observation is:
{observation}
reward: {reward}
done: {done}
Decide your next action(s).
Your reponse should be in the format of <think>...</think><answer>...</answer>
Your response should STRICTLY follow the format:
<think>...</think><answer>...</answer>
"""
 No newline at end of file
+13 −11
Original line number Diff line number Diff line
@@ -19,30 +19,32 @@ python -m vagen.env.sokoban.create_dataset \
    --n_candidate 20000 \
    --force-gen


# max_trajectory_length = max_prompt_length + max_response_length
#Set use_remove_padding to false, if true, causing batch size must be postive error in vllm

python3 -m vagen.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    algorithm.high_level_gamma=0.95 \
    data.train_files=data/sokoban-text-1-step/train.parquet \
    data.val_files=data/sokoban-text-1-step/test.parquet \
    data.train_batch_size=64 \
    data.train_batch_size=512 \
    data.max_prompt_length=768 \
    data.max_response_length=128 \
    data.max_response_length=256 \
    data.max_trajectory_length=1024 \
    data.image_key=images \
    actor_rollout_ref.model.path=Qwen/Qwen2.5-1.5B-Instruct \
    actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.model.use_remove_padding=False \
    actor_rollout_ref.actor.ppo_mini_batch_size=32 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \
    actor_rollout_ref.actor.ppo_mini_batch_size=64 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
    actor_rollout_ref.actor.use_kl_loss=True \
    actor_rollout_ref.actor.kl_loss_coef=0.01 \
    actor_rollout_ref.actor.kl_loss_coef=0.001 \
    actor_rollout_ref.actor.kl_loss_type=mse \
    actor_rollout_ref.model.enable_gradient_checkpointing=True \
    actor_rollout_ref.actor.fsdp_config.param_offload=False \
    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
@@ -50,14 +52,14 @@ python3 -m vagen.trainer.main_ppo \
    actor_rollout_ref.rollout.enforce_eager=False \
    actor_rollout_ref.rollout.free_cache_engine=False \
    actor_rollout_ref.rollout.n=1 \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
    actor_rollout_ref.ref.fsdp_config.param_offload=True \
    +actor_rollout_ref.ref.use_ref=True \
    algorithm.kl_ctrl.kl_coef=0.001 \
    trainer.critic_warmup=0 \
    trainer.logger=['console','wandb'] \
    trainer.project_name='vagen' \
    trainer.experiment_name='debug_single_action_single_turn_grpo' \
    trainer.experiment_name='debug_single_action_single_turn_grpo_0_5B_kl_strict_format' \
    trainer.n_gpus_per_node=1 \
    trainer.nnodes=1 \
    trainer.save_freq=100 \
@@ -67,6 +69,6 @@ python3 -m vagen.trainer.main_ppo \
    rollout_manager.window_size=5 \
    trainer.val_before_train=True \
    trainer.val_generations_to_log_to_wandb=8 \
    rollout_manager.n_trajectory=2 \
    rollout_manager.n_trajectory=1 \
    rollout_manager.use_loss_mask=True \
    2>&1 | tee debug_single_action_single_turn_grpo.log
    2>&1 | tee debug_single_action_single_turn_grpo_0_5B_kl_strict_format.log
+82 −0
Original line number Diff line number Diff line
set -x

export VLLM_ATTENTION_BACKEND=XFORMERS
export PYTHONHASHSEED=0

python -m vagen.env.sokoban.create_dataset \
    --data_dir data/sokoban-text-3-step \
    --max_action_length 3 \
    --dim_room 6 6 \
    --num_boxes 1 \
    --max_steps 100 \
    --search_depth 30 \
    --start_seed 0 \
    --train_ratio 0.8 \
    --max_action_per_step 1 \
    --max_action_penalty 0.0 \
    --format_reward 0.5 \
    --n_candidate 50000 \
    --force-gen

if [ $? -ne 0 ]; then
    echo "Failed to generate dataset"
    exit 1
fi

python3 -m vagen.trainer.main_ppo \
    algorithm.adv_estimator=masked_gae \
    algorithm.high_level_gamma=0.95 \
    data.train_files=data/sokoban-text-3-step/train.parquet \
    data.val_files=data/sokoban-text-3-step/test.parquet \
    data.train_batch_size=128 \
    data.max_prompt_length=512 \
    data.max_response_length=256 \
    data.max_trajectory_length=1664 \
    data.image_key=images \
    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.model.use_remove_padding=False \
    actor_rollout_ref.actor.ppo_mini_batch_size=32 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \
    actor_rollout_ref.actor.use_kl_loss=False \
    actor_rollout_ref.actor.kl_loss_coef=0.001 \
    actor_rollout_ref.actor.kl_loss_type=mse \
    actor_rollout_ref.model.enable_gradient_checkpointing=True \
    actor_rollout_ref.actor.fsdp_config.param_offload=False \
    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \
    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
    actor_rollout_ref.rollout.enable_chunked_prefill=False \
    actor_rollout_ref.rollout.enforce_eager=False \
    actor_rollout_ref.rollout.free_cache_engine=False \
    actor_rollout_ref.rollout.n=1 \
    actor_rollout_ref.rollout.temperature=0.7 \
    actor_rollout_ref.rollout.top_p=0.95 \
    +actor_rollout_ref.ref.use_ref=True \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \
    actor_rollout_ref.ref.fsdp_config.param_offload=True \
    critic.optim.lr=1e-5 \
    critic.model.use_remove_padding=False \
    critic.model.path=Qwen/Qwen2.5-3B-Instruct \
    critic.model.enable_gradient_checkpointing=True \
    critic.ppo_micro_batch_size_per_gpu=1 \
    critic.model.fsdp_config.param_offload=False \
    critic.model.fsdp_config.optimizer_offload=False \
    algorithm.kl_ctrl.kl_coef=0.001 \
    trainer.critic_warmup=0 \
    trainer.logger=['console','wandb'] \
    trainer.project_name='vagen' \
    trainer.experiment_name='debug_single_action_3_turns_ppo_3B_masked_gae_temp_0.7_top_p_0.95' \
    trainer.n_gpus_per_node=4 \
    trainer.nnodes=1 \
    trainer.save_freq=400 \
    trainer.test_freq=5 \
    trainer.total_epochs=15 \
    rollout_manager.max_turns=3 \
    rollout_manager.window_size=5 \
    trainer.val_before_train=True \
    trainer.val_generations_to_log_to_wandb=8 \
    rollout_manager.n_trajectory=1 \
    2>&1 | tee debug_single_action_3_turns_ppo_3B_masked_gae_temp_0.7_top_p_0.95.log
+82 −0
Original line number Diff line number Diff line
set -x

export VLLM_ATTENTION_BACKEND=XFORMERS
export PYTHONHASHSEED=0

python -m vagen.env.sokoban.create_dataset \
    --data_dir data/sokoban-text-3-step \
    --max_action_length 3 \
    --dim_room 6 6 \
    --num_boxes 1 \
    --max_steps 100 \
    --search_depth 30 \
    --start_seed 0 \
    --train_ratio 0.8 \
    --max_action_per_step 1 \
    --max_action_penalty 0.0 \
    --format_reward 0.5 \
    --n_candidate 50000 \
    --force-gen

if [ $? -ne 0 ]; then
    echo "Failed to generate dataset"
    exit 1
fi

python3 -m vagen.trainer.main_ppo \
    algorithm.adv_estimator=multi_turn_gae \
    algorithm.high_level_gamma=0.95 \
    data.train_files=data/sokoban-text-3-step/train.parquet \
    data.val_files=data/sokoban-text-3-step/test.parquet \
    data.train_batch_size=128 \
    data.max_prompt_length=512 \
    data.max_response_length=256 \
    data.max_trajectory_length=1664 \
    data.image_key=images \
    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.model.use_remove_padding=False \
    actor_rollout_ref.actor.ppo_mini_batch_size=32 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \
    actor_rollout_ref.actor.use_kl_loss=False \
    actor_rollout_ref.actor.kl_loss_coef=0.001 \
    actor_rollout_ref.actor.kl_loss_type=mse \
    actor_rollout_ref.model.enable_gradient_checkpointing=True \
    actor_rollout_ref.actor.fsdp_config.param_offload=False \
    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \
    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
    actor_rollout_ref.rollout.enable_chunked_prefill=False \
    actor_rollout_ref.rollout.enforce_eager=False \
    actor_rollout_ref.rollout.free_cache_engine=False \
    actor_rollout_ref.rollout.n=1 \
    actor_rollout_ref.rollout.temperature=0.7 \
    actor_rollout_ref.rollout.top_p=0.95 \
    +actor_rollout_ref.ref.use_ref=True \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \
    actor_rollout_ref.ref.fsdp_config.param_offload=True \
    critic.optim.lr=1e-5 \
    critic.model.use_remove_padding=False \
    critic.model.path=Qwen/Qwen2.5-3B-Instruct \
    critic.model.enable_gradient_checkpointing=True \
    critic.ppo_micro_batch_size_per_gpu=1 \
    critic.model.fsdp_config.param_offload=False \
    critic.model.fsdp_config.optimizer_offload=False \
    algorithm.kl_ctrl.kl_coef=0.001 \
    trainer.critic_warmup=0 \
    trainer.logger=['console','wandb'] \
    trainer.project_name='vagen' \
    trainer.experiment_name='debug_single_action_3_turns_ppo_3B_multi_turn_gae_temp_0.7_top_p_0.95' \
    trainer.n_gpus_per_node=4 \
    trainer.nnodes=1 \
    trainer.save_freq=400 \
    trainer.test_freq=5 \
    trainer.total_epochs=15 \
    rollout_manager.max_turns=3 \
    rollout_manager.window_size=5 \
    trainer.val_before_train=True \
    trainer.val_generations_to_log_to_wandb=8 \
    rollout_manager.n_trajectory=1 \
    2>&1 | tee debug_single_action_3_turns_ppo_3B_multi_turn_gae_temp_0.7_top_p_0.95.log
Loading