update rollout; fix bug

        Args:

            template: The text part of the observation

            contents: A list of content objects to include in the observation

            contents: A list of content objects to include in the observation (str, int, Image.Image, ...)

            replace_keys: The key to replace in the observation template

        This function will:

            assert replace_keys[i] in result_template, f"replace_keys[{i}] must be in text"

            if isinstance(content, str):

                result_template = result_template.replace(replace_keys[i], content)

            if not isinstance(content, Image.Image):

                # transform content to string

                result_template = result_template.replace(replace_keys[i], str(content))

                continue

            # Create the placeholder for this image

                'env_config': self.env_config,

                'seed': seed_idx

            }

            # TODO: no reward model defined here for the reward will be generated while rollout

            return {

                "data_source": self.env_name,

                "prompt": [{"role": "user", "content": instruction}],

                "ability": "bfs",

                "reward_model": {"style": "rule", "ground_truth": {"target": 0, "numbers": [0, 0]}},

                "extra_info": {"split": split, **env_settings}

            }

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument('--config_path', type=str, default='config.yaml')

    parser.add_argument('--start_seed', type=int, default=0)

    parser.add_argument('--train_size', type=int, default=100)

    parser.add_argument('--test_size', type=int, default=100)

    parser.add_argument('--start_seed', type=int, default=10000)

    parser.add_argument('--train_size', type=int, default=10000)

    parser.add_argument('--test_size', type=int, default=1000)

    args = parser.parse_args()

    creator = DatasetCreator(config_path=args.config_path)

    creator.create_dataset(start_seed=args.start_seed, train_size=args.train_size, test_size=args.test_size)

    creator.merge_parquet_files(source_files=[f'{creator.data_dir}/train.parquet', f'{creator.data_dir}/test.parquet'], output_file=f'{creator.data_dir}/merged.parquet')

import argparse

import datasets

from vagen.env.create_dataset import DatasetCreator

class SokobanDatasetCreator(DatasetCreator):

def main():

    # Parse command-line arguments

    parser = argparse.ArgumentParser(description="Generate trajectories using specified environment and policy.")

    parser.add_argument("--start_seed", type=int, default=10000, help="Seed for random number generation (default: 10000).")

    parser.add_argument("--output", type=str, default="data/sokoban", help="Output file to save the trajectories (default: 'data/sokoban').")

    parser.add_argument("--train_size", type=int, default=300, help="Number of trajectories to generate (default: 3000).")

    parser.add_argument("--test_size", type=int, default=10, help="Number of trajectories to generate (default: 100).")

    parser.add_argument("--config", type=str, default="config/sokoban.yaml", help="Config file to use (default: 'config/sokoban.yaml').")

    args = parser.parse_args()

    assert args.env == "sokoban", "Unsupported environment: {args.env}"

    assert args.algo == "bfs", "Unsupported algorithm: {args.algo}"

    data_source = args.env

    dim_x, dim_y, num_boxes, max_steps, search_depth = os.environ.get("DIM_X"), os.environ.get("DIM_Y"), os.environ.get("NUM_BOXES"), os.environ.get("MAX_STEPS"), os.environ.get("SEARCH_DEPTH")

    dim_x, dim_y, num_boxes, max_steps, search_depth = int(dim_x), int(dim_y), int(num_boxes), int(max_steps), int(search_depth)

    seeds = range(args.seed, args.seed + args.train_size + args.test_size)

    instructions = []

    for seed in seeds:

        env = SokobanEnv(

            dim_room=(dim_x, dim_y),

            num_boxes=num_boxes,

            max_steps=max_steps,

            search_depth=search_depth

        )

        observation = env.reset(seed=seed, mode='tiny_rgb_array')

        instruction = INSTRUCTION_TEMPLATE.format(observation=observation)

        instructions.append(instruction)

    def create_filtered_dataset(

        self,

        start_seed: int,

        train_size: int,

        test_size: int,

        max_steps: int = 5

    ):

        # TODO

        return 

    def _create_instance(idx, instruction):

        prompt_formatted = templates[args.prefix].format(prompt=instruction)

        return {

            "data_source": data_source,

            "prompt": [{"role": "user", "content": prompt_formatted}],

            "ability": "bfs",

            "reward_model": {"style": "rule", "ground_truth": {"target": 0, "numbers": [0, 0]}},

            "extra_info": {"split": "train", "index": idx}

        }

    train_dataset = Dataset.from_list([_create_instance(args.seed + i, instructions[i]) for i in range(args.train_size)])

    test_dataset = Dataset.from_list([_create_instance(args.seed + i, instructions[i]) for i in range(args.train_size, args.train_size + args.test_size)])

    def make_map_fn(split):

        def process_fn(example, idx):

            return example

        return process_fn

    train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True)

    test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True)

    train_dataset.to_parquet(os.path.join(args.output, 'train.parquet'))

    test_dataset.to_parquet(os.path.join(args.output, 'test.parquet'))

if __name__ == "__main__":

    main()

 No newline at end of file

    parser = argparse.ArgumentParser()

    parser.add_argument('--config_path', type=str, default='vagen/env/sokoban.yaml')

    parser.add_argument('--start_seed', type=int, default=0)

    parser.add_argument('--train_size', type=int, default=100)

    parser.add_argument('--test_size', type=int, default=100)

    args = parser.parse_args()

    creator = SokobanDatasetCreator(config_path=args.config_path)

    creator.create_filtered_dataset(start_seed=args.start_seed, train_size=args.train_size, test_size=args.test_size)

    algorithm.adv_estimator=grpo \

    data.train_files=data/sokoban/train.parquet \

    data.val_files=data/sokoban/test.parquet \

    data.train_batch_size=8 \

    data.train_batch_size=16 \

    data.max_prompt_length=1024 \

    data.max_response_length=2048 \

    data.image_key=images \

    actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-3B-Instruct \

    actor_rollout_ref.actor.optim.lr=1e-6 \

    actor_rollout_ref.model.use_remove_padding=True \

    actor_rollout_ref.actor.ppo_mini_batch_size=2 \

    actor_rollout_ref.actor.ppo_mini_batch_size=4 \

    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \

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

    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \

    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \

    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \

    actor_rollout_ref.rollout.name=vllm \

    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \

    actor_rollout_ref.rollout.enable_chunked_prefill=False \

    trainer.logger=['console','wandb'] \

    trainer.project_name='vagen' \

    trainer.experiment_name='qwen2_5_vl_3b_function_rm' \

    trainer.n_gpus_per_node=2 \

    trainer.n_gpus_per_node=4 \

    trainer.nnodes=1 \

    trainer.save_freq=-1 \

    trainer.test_freq=5 \