dev minor

from vagen.env.register import REGISTERED_ENVS, register

from vagen.env.sokoban.env import SokobanInterface

 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 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']

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

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

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

    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')

        # 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.")

            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

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

            env_settings = {

                'env_name': self.env_name,

                'env_config': self.env_config,

                'interface_config': self.interface_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": ''}],

                "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)

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

REGISTERED_ENVS = {}

def register(cls=None, *, name=None):

    """

    A decorator to register environment classes in the REGISTERED_ENVS dictionary.

    Args:

        cls: The class to register

        name: Optional custom name for the environment. If not provided, 

              the class name will be used

    Usage:

        @register

        class SokobanEnv(BaseEnv):

            pass

        @register(name="custom_sokoban")

        class CustomSokobanEnv(BaseEnv):

            pass

    """

    def _register(cls):

        # Use provided name or class name as the registry key

        key = name if name is not None else cls.__name__

        # Register the class

        REGISTERED_ENVS[key] = cls

        # Return the class unchanged so it can be used normally

        return cls

    # Handle case when decorator is used with no arguments: @register

    if cls is not None:

        return _register(cls)

    # Handle case when decorator is used with arguments: @register(name="...")

    return _register

 No newline at end of file