minor (d9b66ead) · Commits · 张泽凯 / VAGEN

vagen/env/svg/env.py

+10 −71

Original line number	Diff line number	Diff line
		@@ -42,10 +42,6 @@ class SVGEnv(BaseEnv):
		if hasattr(self.config, "seed") and self.config.seed is not None:
		self.rng.seed(self.config.seed)

		# Set up analysis logging if enabled
		if self.config.analysis_mode:
		self._setup_analysis_logging()


		def reset(self, seed=None) -> Tuple[Dict, Dict]:
		"""Reset the environment with an optional seed"""
		@@ -190,15 +186,7 @@ class SVGEnv(BaseEnv):

		def close(self):
		"""Close the environment and clean up resources"""
		if hasattr(self, 'failure_logger'):
		for handler in self.failure_logger.handlers:
		handler.close()
		self.failure_logger.removeHandler(handler)

		if hasattr(self, 'success_logger'):
		for handler in self.success_logger.handlers:
		handler.close()
		self.success_logger.removeHandler(handler)
		pass

		def _render(self, init_obs=False):
		"""Render the current state of the environment"""
		@@ -236,62 +224,20 @@ class SVGEnv(BaseEnv):
		"multi_modal_data": multi_modal_data,
		}

		def _setup_analysis_logging(self):
		"""Set up logging for analysis mode"""
		log_dir = Path(self.config.data_dir) / 'analysis_logs'
		os.makedirs(log_dir, exist_ok=True)

		# Failure logger
		self.failure_logger = logging.getLogger(f'svg_failure_{id(self)}')
		self.failure_logger.setLevel(logging.INFO)

		if not self.failure_logger.handlers:
		failure_handler = logging.FileHandler(log_dir / 'failure_cases.log')
		failure_handler.setFormatter(logging.Formatter('%(message)s'))
		self.failure_logger.addHandler(failure_handler)

		# Success logger
		self.success_logger = logging.getLogger(f'svg_success_{id(self)}')
		self.success_logger.setLevel(logging.INFO)

		if not self.success_logger.handlers:
		success_handler = logging.FileHandler(log_dir / 'success_cases.log')
		success_handler.setFormatter(logging.Formatter('%(message)s'))
		self.success_logger.addHandler(success_handler)

		def set_dino_model(self, model):
		self.dino_model = model

		if __name__ == "__main__":
		config = SvgEnvConfig(
		dataset_name="starvector/svg-emoji-simple",
		data_dir="vagen/env/svg/data",
		split="test",
		model_size="small"
		)
		config = SvgEnvConfig()

		try:
		env = SVGEnv(config)
		print(f"Successfully loaded dataset")

		# Test with seed
		seed = 42
		obs, info = env.reset(seed=seed)
		print(f"Testing with seed {seed}")

		# Example SVG action
		action = """<think>
		The image appears to be a simple emoji face with two eyes and a smile.
		I'll create an SVG with:
		1. A circle for the face
		2. Two circles for the eyes
		3. A path for the smile
		</think>
		<answer>
		print(env.system_prompt())

		obs, info = env.reset()
		print(obs["obs_str"])

		action = """<answer>
		<svg width="100" height="100" viewBox="0 0 100 100">
		<circle cx="50" cy="50" r="40" fill="yellow"/>
		<circle cx="35" cy="40" r="5" fill="black"/>
		<circle cx="65" cy="40" r="5" fill="black"/>
		<path d="M30 60 Q50 75 70 60" stroke="black" stroke-width="3" fill="none"/>
		</svg>
		</answer>"""
		@@ -299,14 +245,7 @@ if __name__ == "__main__":
		obs, reward, done, info = env.step(action)
		print(f"Reward: {reward}")
		print(f"Done: {done}")
		print(f"obs:{obs}")
		print(f"Score components: {info.get('scores', {})}")

		# Test with another seed to verify determinism
		seed = 123
		obs, info = env.reset(seed=seed)
		print(f"\nTesting with seed {seed}")
		print(obs["obs_str"])

		print(f"Total reward: {env.compute_reward()}")
		env.close()
		No newline at end of file
		except Exception as e:
		print(f"Error: {e}")
		No newline at end of file

vagen/env/svg/env_config.py

+2 −17

Original line number	Diff line number	Diff line
		@@ -15,14 +15,9 @@ class SvgEnvConfig(BaseEnvConfig):
		dino_only: bool = False
		dino_weight: Optional[float] = None
		structural_weight: Optional[float] = None
		color_weight: Optional[float] = None
		code_weight: Optional[float] = None
		# Reward configuration
		format_reward: float = 0.5
		format_penalty: float = 0.0
		# Analysis mode for logging
		analysis_mode: bool = False


		def config_id(self) -> str:
		"""Generate a unique identifier for this configuration"""
		@@ -39,7 +34,7 @@ class SvgEnvConfig(BaseEnvConfig):
		if field.name in id_fields])

		# Add optional fields if they're set
		optional_fields = ["dino_weight", "structural_weight", "color_weight", "code_weight"]
		optional_fields = ["dino_weight", "structural_weight"]
		for field_name in optional_fields:
		value = getattr(self, field_name)
		if value is not None:
		@@ -59,23 +54,13 @@ class SvgEnvConfig(BaseEnvConfig):
		score_config["dino_weight"] = self.dino_weight
		if self.structural_weight is not None:
		score_config["structural_weight"] = self.structural_weight
		if self.color_weight is not None:
		score_config["color_weight"] = self.color_weight
		if self.code_weight is not None:
		score_config["code_weight"] = self.code_weight

		return score_config


		if __name__ == "__main__":
		# Example usage
		config = SvgEnvConfig(
		dataset_name="starvector/svg-emoji-simple",
		data_dir="data/svg",
		model_size="small",
		dino_only=False,
		dino_weight=5.0
		)
		config = SvgEnvConfig()

		print(config.config_id())
		print(config.get_score_config())
		No newline at end of file

vagen/env/svg/service.py

+7 −79

Original line number	Diff line number	Diff line
		@@ -9,7 +9,7 @@ from vagen.env.svg.score import calculate_total_score, calculate_total_score_bat
		from vagen.env.svg.dino import get_dino_model
		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
		import logging
		from PIL import Image

		class SVGService(BaseService):
		"""
		@@ -32,13 +32,11 @@ class SVGService(BaseService):
		self.cache = {}

		# Load the DINO model directly in the service
		# This allows all environments to share the same model instance
		self.model_size = model_size
		self.dino_model = None # Will be loaded on first use

		# Store device for model inference
		self.device = "cuda" if torch.cuda.is_available() else "cpu"
		logging.info(f"SVGService initialized with {max_workers} workers, model_size={model_size}, device={self.device}")

		def _get_dino_model(self):
		"""
		@@ -46,7 +44,6 @@ class SVGService(BaseService):
		Uses lazy loading to avoid loading the model until needed.
		"""
		if self.dino_model is None:
		logging.info(f"Loading DINO model (size={self.model_size}, device={self.device})")
		self.dino_model = get_dino_model(self.model_size, self.device)
		return self.dino_model

		@@ -57,18 +54,12 @@ class SVGService(BaseService):
		Args:
		ids2configs: A dictionary where each key is an environment ID and the corresponding
		value is the configuration for that environment.
		Each config should contain:
		- env_name: Should be "SVG"
		- env_config: SVG specific configuration
		"""
		# Define worker function
		def create_single_env(env_id, config):
		# Verify environment type
		env_name = config.get('env_name', 'svg')
		if env_name != 'svg':
		return env_id, None, f"Expected environment type 'SVG', got '{env_name}'"

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

		@@ -79,8 +70,6 @@ class SVGService(BaseService):
		env = SVGEnv(env_config)

		return env_id, (env, env_config), None
		except Exception as e:
		return env_id, None, str(e)

		# Use ThreadPoolExecutor for parallel creation
		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
		@@ -94,33 +83,26 @@ class SVGService(BaseService):
		for future in as_completed(futures):
		env_id = futures[future]
		env_id, result, error = future.result()
		if error:
		logging.error(f"Error creating environment {env_id}: {error}")
		continue

		if not error:
		env, env_config = result
		self.environments[env_id] = env
		self.env_configs[env_id] = env_config
		# Initialize cache for this environment
		self.cache[env_id] = {}
		logging.info(f"Environment {env_id} created successfully")

		def reset_batch(self, ids2seeds: Dict[Any, Any]) -> Dict[Any, Tuple[Any, Any]]:
		"""
		Reset multiple SVG environments in parallel.

		Args:
		ids2seeds: A dictionary where each key is an environment ID and the corresponding
		value is a seed value (or None for using default seeding behavior).
		ids2seeds: A dictionary mapping environment IDs to seed values.

		Returns:
		A dictionary mapping environment IDs to tuples of the form (observation, info)
		"""
		results = {}

		# Define worker function
		def reset_single_env(env_id, seed):
		try:
		if env_id not in self.environments:
		return env_id, None, f"Environment {env_id} not found"

		@@ -140,8 +122,6 @@ class SVGService(BaseService):
		# Serialize the observation for return
		serialized_observation = serialize_observation(observation)
		return env_id, (serialized_observation, info), None
		except Exception as e:
		return env_id, None, str(e)

		# Use ThreadPoolExecutor for parallel reset
		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
		@@ -156,7 +136,6 @@ class SVGService(BaseService):
		env_id = futures[future]
		env_id, result, error = future.result()
		if error:
		logging.error(f"Error resetting environment {env_id}: {error}")
		results[env_id] = ({}, {"error": error})
		else:
		results[env_id] = result
		@@ -288,16 +267,12 @@ class SVGService(BaseService):
		"""
		results = {}

		# Define worker function
		def compute_reward_single_env(env_id):
		try:
		if env_id not in self.environments:
		return env_id, None, f"Environment {env_id} not found"

		env = self.environments[env_id]
		return env_id, env.compute_reward(), None
		except Exception as e:
		return env_id, None, str(e)

		# Use ThreadPoolExecutor for parallel computation
		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
		@@ -312,7 +287,6 @@ class SVGService(BaseService):
		env_id = futures[future]
		env_id, result, error = future.result()
		if error:
		logging.error(f"Error computing reward for environment {env_id}: {error}")
		results[env_id] = 0.0
		else:
		results[env_id] = result
		@@ -331,16 +305,12 @@ class SVGService(BaseService):
		"""
		results = {}

		# Define worker function
		def get_system_prompt_single_env(env_id):
		try:
		if env_id not in self.environments:
		return env_id, None, f"Environment {env_id} not found"

		env = self.environments[env_id]
		return env_id, env.system_prompt(), None
		except Exception as e:
		return env_id, None, str(e)

		# Use ThreadPoolExecutor for parallel retrieval
		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
		@@ -355,7 +325,6 @@ class SVGService(BaseService):
		env_id = futures[future]
		env_id, result, error = future.result()
		if error:
		logging.error(f"Error getting system prompt for environment {env_id}: {error}")
		results[env_id] = ""
		else:
		results[env_id] = result
		@@ -372,20 +341,14 @@ class SVGService(BaseService):
		# If no env_ids provided, close all environments
		if env_ids is None:
		env_ids = list(self.environments.keys())
		logging.info(f"Environments {env_id} needs to close")


		# Define worker function
		def close_single_env(env_id):
		try:
		if env_id not in self.environments:
		return f"Environment {env_id} not found"

		env = self.environments[env_id]
		env.close()
		return None
		except Exception as e:
		return str(e)

		# Use ThreadPoolExecutor for parallel closing
		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
		@@ -394,16 +357,13 @@ class SVGService(BaseService):

		# Wait for all tasks to complete
		for future in as_completed(futures):
		error = future.result()
		if error:
		logging.error(f"Error closing environment: {error}")
		future.result()

		# Remove closed environments from dictionaries
		for env_id in env_ids:
		self.environments.pop(env_id, None)
		self.env_configs.pop(env_id, None)
		self.cache.pop(env_id, None)
		logging.info(f"Environment {env_id} closed successfully")

		def _process_svg_actions_batch(self, ids2actions):
		"""
		@@ -421,7 +381,6 @@ class SVGService(BaseService):

		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
		def process_action(env_id, action):
		try:
		if env_id not in self.environments:
		return env_id, None, f"Environment {env_id} not found"

		@@ -482,7 +441,7 @@ class SVGService(BaseService):
		"metrics": metrics,
		"info": info,
		"valid": False,
		"done": False # Changed from True to False to allow training to continue
		"done": False
		}, None

		# Process SVG (valid or invalid)
		@@ -491,15 +450,7 @@ class SVGService(BaseService):
		env.gen_svg_code = rst['actions'][0]
		env.valid_actions = rst['actions']

		try:
		# Try to rasterize SVG, use a blank image as fallback
		try:
		_, gen_image = process_and_rasterize_svg(env.gen_svg_code)
		env.gen_image = gen_image
		except Exception as e:
		logging.error(f"Error rasterizing SVG for {env_id}: {e}")
		# Create a blank white image as fallback
		env.gen_image = Image.new('RGB', (256, 256), color='white')
		_, env.gen_image = process_and_rasterize_svg(env.gen_svg_code)

		return env_id, {
		"env": env,
		@@ -507,31 +458,10 @@ class SVGService(BaseService):
		"gen_svg_code": env.gen_svg_code,
		"rst": rst,
		"metrics": metrics,
		"valid": True, # Consider it valid for processing even if SVG is invalid
		"valid": True,
		"done": False
		}, None

		except Exception as e:
		logging.error(f"Error in SVG processing pipeline for {env_id}: {e}")
		# Even if processing fails, don't mark as done
		info = rst.copy()
		info["metrics"] = metrics
		info["error"] = str(e)
		return env_id, {
		"env": env,
		"gen_image": None,
		"gen_svg_code": env.gen_svg_code,
		"rst": rst,
		"metrics": metrics,
		"info": info,
		"valid": False,
		"done": False # Changed from True to False to allow training to continue
		}, None

		except Exception as e:
		logging.error(f"Unexpected error in action processing for {env_id}: {e}")
		return env_id, None, str(e)

		# Submit all action processing tasks
		futures = {
		executor.submit(process_action, env_id, action): env_id
		@@ -543,8 +473,6 @@ class SVGService(BaseService):
		env_id = futures[future]
		env_id, result, error = future.result()
		if error:
		logging.error(f"Error processing action for environment {env_id}: {error}")
		# Return error but don't mark as done
		error_results[env_id] = ({}, 0.0, False, {"error": error})
		else:
		env_processing_results[env_id] = result

vagen/env/svg/svg_utils.py

+1 −0

Original line number	Diff line number	Diff line
		@@ -99,6 +99,7 @@ def rasterize_svg(svg_string, resolution=224, dpi = 128, scale=2):



		# -------------- download/pre-load dataset --------------

		def load_svg_dataset(data_dir, dataset_name, split):
		"""Load the SVG dataset from local files or HuggingFace"""

Admin message