minor update

        )

        return f"{env_config_str}, {interface_config_str}"

    def get_task_instruction(self) -> str:

        return instruction_template

        return instruction_template.format(

            max_action_per_step=self.interface_config['max_action_per_step']

        )

    def get_traj_reward(self):

        return self.traj_reward

1. Push boxes (can't pull).

2. Avoid walls (#).

Actions you can take: Up, Down, Left, Right. You can only take one action at a time.

Actions you can take: Up, Down, Left, Right. You can only take up to {max_action_per_step} 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).

Please think step by step and provide the action you want to take.

Your reponse should be in the format of <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>

"""

# 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

# If the reponse provides a final answer and is corect, the model receives an accurtacy reward of +1

# is the response encloses its thinking in <think></think> and the final answer is <answer></answer> tags, the model receives a format reward of +1

    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.use_ref=False \

    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \

    actor_rollout_ref.ref.fsdp_config.param_offload=True \

    critic.optim.lr=1e-5 \

set -x

export VLLM_ATTENTION_BACKEND=XFORMERS

export PYTHONHASHSEED=0

python -m vagen.env.sokoban.create_dataset \

    --data_dir data/sokoban-text-1-step \

    --max_action_length 1 \

    --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.1 \

    --format_reward 0.5 \

    --n_candidate 20000 \

    --force-gen

if [ $? -ne 0 ]; then

    echo "Failed to generate dataset"

    exit 1

fi

# max_trajectory_length = max_prompt_length + max_response_length

python3 -m vagen.trainer.main_ppo \

    algorithm.adv_estimator=gae \

    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=16 \

    data.max_prompt_length=768 \

    data.max_response_length=128 \

    data.max_trajectory_length=1024 \

    data.image_key=images \

    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=4 \

    actor_rollout_ref.actor.use_kl_loss=True \

    actor_rollout_ref.actor.kl_loss_coef=0.01 \

    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=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 \

    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.ref.use_ref=True \

    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \

    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-0.5B-Instruct \

    critic.model.enable_gradient_checkpointing=True \

    critic.ppo_micro_batch_size_per_gpu=4 \

    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_single_turn_ppo_kl' \

    trainer.n_gpus_per_node=1 \

    trainer.nnodes=1 \

    trainer.save_freq=100 \

    trainer.test_freq=5 \

    trainer.total_epochs=15 \

    rollout_manager.max_turns=1 \

    rollout_manager.window_size=5 \

    trainer.val_before_train=True \

    trainer.val_generations_to_log_to_wandb=8 \

    rollout_manager.n_trajectory=8 \

    2>&1 | tee debug_single_action_single_turn_ppo_kl.log

        Role.Critic: ray.remote(CriticWorker),

    }

    use_ref = (config.algorithm.adv_estimator not in  [AdvantageEstimator.GAE, AdvantageEstimator.MULTI_TURN_GAE] and \

        config.actor_rollout_ref.ref.get('use_ref', True))

    # use_ref = (config.algorithm.adv_estimator not in  [AdvantageEstimator.GAE, AdvantageEstimator.MULTI_TURN_GAE] and \

    #     config.actor_rollout_ref.ref.get('use_ref', True))

    use_ref = config.actor_rollout_ref.ref.get('use_ref', True)

    print(f"[DEBUG] use_ref={use_ref}")

    if use_ref:

        role_worker_mapping[Role.RefPolicy] = ray.remote(ActorRolloutRefWorker)