updated svg (8286638e) · Commits · 张泽凯 / VAGEN

scripts/examples/crossview_grpo/run.sh

+1 −1

Original line number	Diff line number	Diff line
		@@ -18,7 +18,7 @@ python3 -m vagen.trainer.main_ppo \
		algorithm.high_level_gamma=0.95 \
		data.train_files=data/crossview/train.parquet \
		data.val_files=data/crossview/test.parquet \
		data.train_batch_size=16 \
		data.train_batch_size=64 \
		data.max_prompt_length=1024 \
		data.max_response_length=648 \
		data.max_trajectory_length=3600 \

vagen/env/svg/env.py

+3 −8

Original line number	Diff line number	Diff line
		@@ -28,7 +28,7 @@ class SVGEnv(BaseEnv):
		reproduces the image as accurately as possible.
		"""

		def __init__(self, config: SvgEnvConfig):
		def __init__(self, config: SvgEnvConfig,dataset):
		"""Initialize the SVG environment.

		Args:
		@@ -36,13 +36,8 @@ class SVGEnv(BaseEnv):
		"""
		BaseEnv.__init__(self)
		self.config = config
		self.script_dir = os.path.dirname(os.path.abspath(__file__))

		# Load the actual SVG dataset
		self.dataset = load_svg_dataset(
		data_dir=os.path.join(self.script_dir,self.config.get("data_dir", "")),
		dataset_name=self.config.dataset_name,
		split=self.config.get("split", "train")
		)

		# Initialize state variables
		self.total_reward = 0
		@@ -55,7 +50,7 @@ class SVGEnv(BaseEnv):
		self.gen_svg_code = None
		self.gen_image = None
		self.dino_model = None

		self.dataset = dataset
		# Store the format prompt function for later use
		self.format_prompt_func = format_prompt[self.config.get('prompt_format', 'free_think')]

vagen/env/svg/service.py

+33 −16

Original line number	Diff line number	Diff line
		@@ -9,7 +9,8 @@ from vagen.env.svg.score import calculate_total_score, calculate_total_score_bat
		from vagen.env.svg.svg_utils import process_and_rasterize_svg, is_valid_svg
		from vagen.env.utils.context_utils import parse_llm_raw_response, convert_numpy_to_PIL
		from .service_config import SVGServiceConfig

		from vagen.env.svg.svg_utils import (process_and_rasterize_svg, is_valid_svg, load_svg_dataset)
		import os
		class SVGService(BaseService):
		"""
		Service class for SVG environments.
		@@ -30,6 +31,9 @@ class SVGService(BaseService):
		self.environments = {}
		self.env_configs = {}
		self.cache = {}
		self.script_dir = os.path.dirname(os.path.abspath(__file__))
		self.dataset = {}


		# Load the DINO model directly in the service
		# This allows all environments to share the same model instance
		@@ -43,6 +47,17 @@ class SVGService(BaseService):
		self.device = "cuda" if torch.cuda.is_available() else "cpu"
		print(f"SVGService initialized with {self.max_workers} workers, model_size={self.model_size}, device={self.device}")

		def _config_to_env_config(self, config):
		env_config_dict = config.get('env_config', {})
		env_config = SvgEnvConfig(**env_config_dict)
		data_dir= os.path.join(self.script_dir, self.config.get("data_dir", ""))
		dataset_name = env_config.dataset_name
		split = env_config.get("split", "train")
		return {
		"dataset_id":"-".join([str(data_dir), str(dataset_name), str(split)]),
		"config": env_config,
		}

		def create_environments_batch(self, ids2configs: Dict[Any, Any]) -> None:
		"""
		Create multiple SVG environments in parallel.
		@@ -51,19 +66,21 @@ class SVGService(BaseService):
		ids2configs: A dictionary where each key is an environment ID and the corresponding
		value is the configuration for that environment.
		"""
		def create_single_env(env_id, config):
		env_name = config.get('env_name', 'svg')
		if env_name != 'svg':
		return env_id, None, f"Expected environment type 'SVG', got '{env_name}'"

		# Get SVG specific configuration
		env_config_dict = config.get('env_config', {})

		# Create environment config
		env_config = SvgEnvConfig(**env_config_dict)
		id_to_env_config = {}
		for env_id, config in ids2configs.items():
		rst=self._config_to_env_config(config)
		dataset_id=rst["dataset_id"]
		env_config=rst["config"]
		if dataset_id not in self.dataset:
		self.dataset[dataset_id]=load_svg_dataset(
		data_dir=os.path.join(self.script_dir,env_config.get("data_dir", "")),
		dataset_name=env_config.dataset_name,
		split=env_config.get("split", "train")
		)
		id_to_env_config[env_id] = (env_config,dataset_id)

		# Create environment
		env = SVGEnv(env_config)
		def create_single_env(env_id, env_config,dataset):
		env = SVGEnv(env_config,dataset)

		return env_id, (env, env_config), None

		@@ -71,8 +88,8 @@ class SVGService(BaseService):
		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
		# Submit all environment creation tasks
		futures = {
		executor.submit(create_single_env, env_id, config): env_id
		for env_id, config in ids2configs.items()
		executor.submit(create_single_env, k, v[0],self.dataset[v[1]]): env_id
		for k, v in id_to_env_config.items()
		}

		# Process results as they complete

Admin message