Loading vagen/env_new/__init__.py +5 −0 Original line number Diff line number Diff line from .sokoban import SokobanEnv,SokobanConfig from .frozenlake import FrozenLakeEnv,FrozenLakeConfig REGISTERED_ENV = { "sokoban": { "env": SokobanEnv, "config": SokobanConfig, }, "frozenlake": { "env": FrozenLakeEnv, "config": FrozenLakeConfig, } } No newline at end of file vagen/env_new/create_dataset.py +121 −95 Original line number Diff line number Diff line import os from vagen.env_new 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: from vagen.env_new.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. def __init__(self, config: Dict): self.config = config self.data_dir = self.config['data_dir'] self.interface_config = self.config['interface_config'] assert "env_name" in self.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) 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) def _create_instance(seed_idx, split: str = 'train'): 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 = { '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.interface_config["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_file_path) test_dataset.to_parquet(test_file_path) print(f"Dataset successfully generated at {self.data_dir}") train_dataset.to_parquet(train_path) print(f"Train dataset with {len(train_instances)} instances saved to {train_path}") 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": 3 }, "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 vagen/env_new/frozenlake/__init__.py 0 → 100644 +2 −0 Original line number Diff line number Diff line from .config import FrozenLakeConfig from .env import FrozenLakeEnv No newline at end of file Loading
vagen/env_new/__init__.py +5 −0 Original line number Diff line number Diff line from .sokoban import SokobanEnv,SokobanConfig from .frozenlake import FrozenLakeEnv,FrozenLakeConfig REGISTERED_ENV = { "sokoban": { "env": SokobanEnv, "config": SokobanConfig, }, "frozenlake": { "env": FrozenLakeEnv, "config": FrozenLakeConfig, } } No newline at end of file
vagen/env_new/create_dataset.py +121 −95 Original line number Diff line number Diff line import os from vagen.env_new 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: from vagen.env_new.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. def __init__(self, config: Dict): self.config = config self.data_dir = self.config['data_dir'] self.interface_config = self.config['interface_config'] assert "env_name" in self.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) 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) def _create_instance(seed_idx, split: str = 'train'): 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 = { '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.interface_config["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_file_path) test_dataset.to_parquet(test_file_path) print(f"Dataset successfully generated at {self.data_dir}") train_dataset.to_parquet(train_path) print(f"Train dataset with {len(train_instances)} instances saved to {train_path}") 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": 3 }, "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
vagen/env_new/frozenlake/__init__.py 0 → 100644 +2 −0 Original line number Diff line number Diff line from .config import FrozenLakeConfig from .env import FrozenLakeEnv No newline at end of file
