updated training for new env

from vagen.env.register import REGISTERED_ENVS, register

from vagen.env.sokoban.env import SokobanInterface

 No newline at end of file

from .sokoban import SokobanEnv,SokobanConfig

from .frozenlake import FrozenLakeEnv,FrozenLakeConfig

REGISTERED_ENV = {

    "sokoban": {

        "env_cls": SokobanEnv,

        "config_cls": SokobanConfig,

    },

    "frozenlake": {

        "env_cls": FrozenLakeEnv,

        "config_cls": FrozenLakeConfig,

    }

}

 No newline at end of file

from abc import ABC, abstractmethod

import re

from typing import Optional, List, Tuple, Any, Dict

from copy import deepcopy

from transformers import AutoTokenizer

import torch

from PIL import Image

import numpy as np

from dataclasses import dataclass, field

IMAGE_PLACEHOLDER = "<image>"

@dataclass

class EnvConfig:

    """

    Dataclass for managing environment configuration.

    """

    env_name: str

    env_config: Dict[str, Any]

    interface_config: Dict[str, Any]

    seed: int

class BaseEnv(ABC):

    @abstractmethod

    def _reset(self, seed: Optional[int] = None) -> Any:

        """

        Reset the environment.

        NOTE: the environment should be same for the same seed

        Args:

            seed: Seed for the environment

        Returns:

            rendered environment

        """

        pass

    @abstractmethod

    def _step(self, action) -> Tuple[Any, float, bool, Dict]:

        """

        Execute one step in the environment.

        NOTE should also handle predefined invalid action (0)

        Args:

            action: Action to take, must be in action space, or default invalid action

        Returns:

            observation (rendered environment), reward, done, info

        """

        pass

    @abstractmethod

    def close(self):

        """Close the environment."""

        pass

    def step(self, action:Any) -> Tuple[Any, Any, Any, Any]:

        """

        Execute one step in the environment.

        Args:

            action: Action to take, must be in action space, or default invalid action

        Returns:

            observation (rendered environment), reward, done, info

        """

        obs,reward,done,info = self._step(action)

        return obs, reward, done, info

    def reset(self, seed: Optional[int] = None) -> Any:

        """

        Reset the environment.

        NOTE: the environment should be same for the same seed

        Args:

            seed: Seed for the environment  

        Returns:

            obs,info

        """

        obs,info = self._reset(seed)

        return obs,info

class BaseInterface(ABC):

    def __init__(self, env_config: Dict, interface_config: Dict = None):

        self.env_config = env_config

        self.interface_config = interface_config

    @classmethod

    def name_repr(cls) -> str:

        """Get the name of the environment."""

        return cls.__name__

    @abstractmethod

    def _reset(self, seed: Optional[int] = None) -> Tuple[Any, float, bool, Dict]:

        """Reset the environment."""

        pass

    @abstractmethod

    def _step(self, action:str) -> Tuple[Any, float, bool, Dict]:

        """Execute action string in the environment."""

        # return observation, reward, done, info

        # info must contain "llm_raw_response" key, which is a string

        pass

    @classmethod

    @abstractmethod

    def config_repr(cls, config: Dict) -> str:

        """Get the config of the environment."""

        pass

    @abstractmethod

    def close(self):

        """Close the environment."""

        pass

    @abstractmethod

    def get_task_instruction(self) -> str:

        """Get the task instruction."""

        pass

    def step(self, action: str) -> Tuple[Dict, float, bool, Dict]:

        """Execute action string in the environment."""

        """Please use the following assertions to validate the output, 

        then you can rewrite the step in your own class to improve the performance"""

        assert isinstance(action, str), f"action must be str, got {type(action)}"

        obs,reward,done,info = self._step(action)

        assert isinstance(reward, (int, float)), f"reward must be int or float, got {type(reward)}"

        assert isinstance(done, bool), f"done must be bool, got {type(done)}"

        assert isinstance(info, dict), f"info must be dict, got {type(info)}"

        assert isinstance(obs, dict), f"obs must be dict, got {type(obs)}"

        assert "llm_raw_response" in info, f"info must contain 'llm_raw_response' key"

        assert isinstance(info["llm_raw_response"], str), f"info['llm_raw_response'] must be str, got {type(info['llm_raw_response'])}"

        assert "text_template" in obs, f"obs must contain 'text_template' key"

        assert isinstance(obs["text_template"], str), f"obs['text_template'] must be str, got {type(obs['text_template'])}"

        if "multi_modal_data" in obs:

            if IMAGE_PLACEHOLDER in obs["multi_modal_data"]:

                assert isinstance(obs["multi_modal_data"][IMAGE_PLACEHOLDER], list), f"obs['multi_modal_data']['<image>'] must be list, got {type(obs['multi_modal_data'][IMAGE_PLACEHOLDER])}"

                for image in obs["multi_modal_data"][IMAGE_PLACEHOLDER]:

                    assert isinstance(image, Image.Image), f"image must be PIL.Image.Image, got {type(image)}"

                len_of_images = len(obs["multi_modal_data"][IMAGE_PLACEHOLDER])

                len_of_image_in_text_template = len(re.findall(IMAGE_PLACEHOLDER, obs["text_template"]))

                assert len_of_images == len_of_image_in_text_template, f"len_of_images must be equal to len_of_image_in_text_template, got {len_of_images} and {len_of_image_in_text_template}"

        return obs, reward, done, info

    def reset(self, seed: int):

        """Reset the environment."""

        assert isinstance(seed, int), f"seed must be int, got {type(seed)}"

        obs, info = self._reset(seed)

        assert isinstance(info, dict), f"info must be dict, got {type(info)}"

        assert isinstance(obs, dict), f"obs must be dict, got {type(obs)}"

        assert "text_template" in obs, f"obs must contain 'text_template' key"

        assert isinstance(obs["text_template"], str), f"obs['text_template'] must be str, got {type(obs['text_template'])}"

        if "multi_modal_data" in obs:

            if IMAGE_PLACEHOLDER in obs["multi_modal_data"]:

                assert isinstance(obs["multi_modal_data"][IMAGE_PLACEHOLDER], list), f"obs['multi_modal_data']['<image>'] must be list, got {type(obs['multi_modal_data'][IMAGE_PLACEHOLDER])}"

                for image in obs["multi_modal_data"][IMAGE_PLACEHOLDER]:

                    assert isinstance(image, Image.Image), f"image must be PIL.Image.Image, got {type(image)}"

                len_of_images = len(obs["multi_modal_data"][IMAGE_PLACEHOLDER])

                len_of_image_in_text_template = len(re.findall(IMAGE_PLACEHOLDER, obs["text_template"]))

                assert len_of_images == len_of_image_in_text_template, f"len_of_images must be equal to len_of_image_in_text_template, got {len_of_images} and {len_of_image_in_text_template}"

        return obs, info

    @abstractmethod

    def get_traj_reward(self) -> float:

        """Get the reward of the environment."""

import os

from vagen.env import REGISTERED_ENV

import numpy as np

import yaml

from datasets import Dataset, load_dataset

import os

import pandas as pd

import argparse

from pathlib import Path

from typing import Union, List, Dict

class DatasetCreator:

    def __init__(self, config: Dict):

        self.config = config

        self.data_dir = self.config['data_dir']

from vagen.env.utils.env_utils import permanent_seed

def create_dataset_from_yaml(yaml_file_path: str, force_gen=False):

    """

    Create dataset from a YAML configuration file.

        self.env_name = self.config['name']

        self.env_config = self.config['env_config']

        self.interface_config = self.config['interface_config']

    Args:

        yaml_file_path (str): Path to the YAML configuration file

        force_gen (bool): Whether to force regeneration of existing datasets

    The YAML file should have the following structure:

    ```

    seed: 42

    train_path: path/to/train.parquet

    test_path: path/to/test.parquet

    env1:

        env_name: sokoban  # or frozenlake

        env_config:

            # parameters to override the default env config

        split: train  # or test

        size: 100  # number of instances

    env2:

        env_name: frozenlake

        env_config:

            # parameters to override the default env config

        split: test

        size: 50

    ```

    If the environment config class (e.g., SokobanConfig, FrozenLakeConfig) has a 

    generate_seeds(size) method, it will be used to generate seeds for that environment.

    """

    if isinstance(yaml_file_path, str):

        with open(yaml_file_path, 'r') as f:

            yaml_config = yaml.safe_load(f)

    else:

        yaml_config = yaml_file_path

    train_path = yaml_config.get('train_path')

    test_path = yaml_config.get('test_path')

    def create_dataset(self, seed: Union[int, List[int]], train_size, test_size, force_gen=False):

        train_file_path = os.path.join(self.data_dir, 'train.parquet')

        test_file_path = os.path.join(self.data_dir, 'test.parquet')

    os.makedirs(os.path.dirname(train_path), exist_ok=True)

    os.makedirs(os.path.dirname(test_path), exist_ok=True)

        # Check if files already exist and force_gen is False

        if not force_gen and os.path.exists(train_file_path) and os.path.exists(test_file_path):

            print(f"Dataset files already exist at {self.data_dir}. Skipping generation.")

    if not force_gen and os.path.exists(train_path) and os.path.exists(test_path):

        print(f"Dataset files already exist at {train_path} and {test_path}. Skipping generation.")

        print(f"Use --force-gen to override and regenerate the dataset.")

        return

        # Ensure data directory exists

        os.makedirs(self.data_dir, exist_ok=True)

        if isinstance(seed, int):

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

        else:

            seeds = seed

    train_instances = []

    test_instances = []

    global_seed = yaml_config.get('seed', 42)

    permanent_seed(global_seed)

        def _create_instance(seed_idx, split: str = 'train'):

    for key, value in yaml_config.items():

        if key in ['train_path', 'test_path','seed']:

            continue

        env_name = value.get('env_name')

        custom_env_config = value.get('env_config', {})

        split = value.get('split', 'train')

        env_size = value.get('size', 100)

        env_config = REGISTERED_ENV[env_name]["config"](**custom_env_config)

        seeds_for_env = None

        if hasattr(env_config, 'generate_seeds'):

            seeds_for_env = env_config.generate_seeds(env_size)

            print(f"Using {len(seeds_for_env)} seeds generated by {env_name} config's generate_seeds method")

        else:

            seeds_for_env = np.random.randint(0, 2**31 - 1, size=env_size).tolist()

        for seed in seeds_for_env:

            env_settings = {

                'env_name': self.env_name,

                'env_config': self.env_config,

                'interface_config': self.interface_config,

                'seed': seed_idx

                'env_name': env_name,

                'env_config': custom_env_config,

                'seed': seed

            }

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

            return {

                "data_source": self.env_name,

            instance = {

                "data_source": env_name,

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

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

            }

        train_instances = [_create_instance(seeds[i], split='train') for i in range(train_size)]

        test_instances = [_create_instance(seeds[train_size + i], split='test') for i in range(test_size)]

        train_dataset = Dataset.from_list(train_instances)

        test_dataset = Dataset.from_list(test_instances)

            if split == 'train':

                train_instances.append(instance)

            else:

                test_instances.append(instance)

    def make_map_fn(split):

        def process_fn(example, idx):

            return example

        return process_fn

    # Create datasets

    if train_instances:

        train_dataset = Dataset.from_list(train_instances)

        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(train_path)

        print(f"Train dataset with {len(train_instances)} instances saved to {train_path}")

        train_dataset.to_parquet(train_file_path)

        test_dataset.to_parquet(test_file_path)

        print(f"Dataset successfully generated at {self.data_dir}")

    def merge_parquet_files(

        self,

        source_files: list[str],

        output_file: str,

        columns: list[str] = None

    ):

        """

        Merge multiple parquet files into a single parquet file.

        Args:

            source_files (list): List of paths to parquet files to merge

            output_file (str): Path to save the merged parquet file

            columns (list, optional): List of columns to include in the merged file.

                                    If None, all columns are included.

        Returns:

            bool: True if successful, False otherwise

        """

        try:

            # Make sure the output directory exists

            output_path = Path(output_file)

            output_path.parent.mkdir(parents=True, exist_ok=True)

            # Initialize an empty DataFrame to hold the merged data

            merged_df = pd.DataFrame()

            for file_path in source_files:

                # Check if the file exists

                if not os.path.exists(file_path):

                    print(f"Warning: File {file_path} does not exist, skipping.")

                    continue

                # Read the parquet file

                df = pd.read_parquet(file_path, columns=columns)

                # Append to the merged DataFrame

                merged_df = pd.concat([merged_df, df], ignore_index=True)

            if merged_df.empty:

                print("No data to merge. Check if source files exist and contain data.")

                return False

            # Write the merged DataFrame to a parquet file

            merged_df.to_parquet(output_file, index=False)

            print(f"Successfully merged {len(source_files)} files into {output_file}")

            return True

        except Exception as e:

            print(f"Error merging parquet files: {str(e)}")

            return False

 No newline at end of file

    if test_instances:

        test_dataset = Dataset.from_list(test_instances)

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

        test_dataset.to_parquet(test_path)

        print(f"Test dataset with {len(test_instances)} instances saved to {test_path}")

    if not train_instances and not test_instances:

        print("No instances were generated. Check your YAML configuration.")

if __name__ == "__main__":

    yaml_file_path = {

        "seed": 42,

        "train_path": "./train_example.parquet",

        "test_path": "./test_example.parquet",

        "env1": {

            "env_name": "sokoban",

            "env_config": {

                "num_boxes": 1

            },

            "split": "train",

            "size": 2

        },

        "env2": {

            "env_name": "frozenlake",

            "env_config": {

                "is_slippery": False,

                "p":0.1

            },

            "split": "test",

            "size": 2

        }

    }

    create_dataset_from_yaml(yaml_file_path, force_gen=True)

    # load the dataset and print

    train_dataset = load_dataset('parquet', data_files={"train": "./train_example.parquet"}, split="train")

    test_dataset = load_dataset('parquet', data_files={"test": "./test_example.parquet"}, split="test")

    for i in range(2):

        print(train_dataset[i])

        print(test_dataset[i])

 No newline at end of file