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

vagen/env/svg/config.py

+2 −2

Original line number	Diff line number	Diff line
		@@ -5,9 +5,9 @@ from typing import Optional, List, Union, Dict
		@dataclass
		class SVGConfig(BaseConfig):
		"""Configuration for the SVG environment"""
		dataset_name: str = "starvector/svg-emoji-simple"
		dataset_name: str = "starvector/svg-icons-simple"
		data_dir: str = "vagen/env/svg/data"
		seed: int = 16
		seed: int = 42
		split: str = "train"
		action_sep: str = "~~"
		# Score configuration

vagen/env/svg/dino.py

+6 −13

Original line number	Diff line number	Diff line
		@@ -13,25 +13,18 @@ from typing import Dict, List, Tuple, Optional, Any, Union

		_model_cache = {}
		_model_cache_lock = threading.Lock()
		_model_counter = 0

		def get_dino_model(model_size="small", device="cuda"):
		"""
		Get or create a DINO model instance with singleton pattern

		Args:
		model_size: Size of DINO model ('small', 'base', or 'large')
		device: Device to run model on ('cuda' or 'cpu')

		Returns:
		DINOScoreCalculator instance
		"""
		# Use the actual DINOScoreCalculator implementation from your code
		# with added singleton pattern
		global _model_counter
		cache_key = f"{model_size}_{device}"

		with _model_cache_lock:
		if cache_key not in _model_cache:
		logging.info(f"Creating new DINO model: {model_size} on {device}")
		_model_counter += 1
		import os
		pid = os.getpid()
		logging.info(f"Process {pid}: Created DINO model #{_model_counter}: {model_size} on {device}")
		_model_cache[cache_key] = DINOScoreCalculator(model_size=model_size, device=device)
		return _model_cache[cache_key]

vagen/env/svg/service.py

+80 −18

Original line number	Diff line number	Diff line
		@@ -231,6 +231,49 @@ class SVGService(BaseService):
		observation = env._render(init_obs=False)
		results[env_id] = serialize_step_result((observation, env.reward, False, info))

		# Step 4: Process invalid or failed generations
		for env_id, result in env_processing_results.items():
		if env_id not in results:
		env = result["env"]

		info = result["rst"].copy() if "rst" in result else {}

		if "metrics" not in info:
		info["metrics"] = {"turn_metrics": {}, "traj_metrics": {}}
		elif "turn_metrics" not in info["metrics"]:
		info["metrics"]["turn_metrics"] = {}
		elif "traj_metrics" not in info["metrics"]:
		info["metrics"]["traj_metrics"] = {}

		info["metrics"]["turn_metrics"]["action_is_valid"] = False
		info["metrics"]["turn_metrics"]["action_is_effective"] = False

		if "scores" not in info:
		info["scores"] = {
		"dino_score": 0.0,
		"structural_score": 0.0,
		"total_score": 0.0
		}

		reward = 0.0

		if hasattr(env.config, "format_penalty"):
		reward = env.config.format_penalty

		env.reward = reward
		env.total_reward += reward
		env.gen_svg_code = None
		env.gen_image = None

		observation = env._render(init_obs=False)

		if env_id in self.cache:
		self.cache[env_id]['gen_image'] = None
		self.cache[env_id]['gen_svg_code'] = None
		self.cache[env_id]['scores'] = info["scores"]

		results[env_id] = serialize_step_result((observation, reward, False, info))

		return results

		def compute_reward_batch(self, env_ids: List[str]) -> Dict[Any, float]:
		@@ -329,6 +372,8 @@ 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):
		@@ -358,6 +403,7 @@ class SVGService(BaseService):
		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):
		"""
		@@ -387,23 +433,32 @@ class SVGService(BaseService):
		action_sep=env.config.get("action_sep", ","),
		max_actions=env.config.get("max_actions_per_step", 1)
		)

		# Handle SVG code extraction
		svg_code = None
		svg_is_valid = False

		# First, try to extract SVG code from the response
		if not rst['actions']:
		svg_code = env._extract_svg_code(action)
		if svg_code and is_valid_svg(svg_code):
		if svg_code:
		svg_is_valid = is_valid_svg(svg_code)
		# Even if SVG is invalid, still keep it for training purposes
		rst['actions'] = [svg_code]
		else:
		svg_code = env._extract_svg_code(rst['actions'][0])
		if svg_code and is_valid_svg(svg_code):
		if svg_code:
		svg_is_valid = is_valid_svg(svg_code)
		# Always keep extracted SVG code regardless of validity
		rst['actions'] = [svg_code]
		else:
		rst['actions'] = []

		# Initialize metrics
		# Initialize metrics - track validity separately from action presence
		metrics = {
		"turn_metrics": {
		"action_is_valid": rst['actions'] != [],
		"action_is_valid": rst['actions'] != [], # Action exists
		"svg_is_valid": svg_is_valid, # SVG syntax is valid
		"action_is_effective": False,
		},
		"traj_metrics": {
		@@ -411,7 +466,7 @@ class SVGService(BaseService):
		}
		}

		# Handle invalid SVG
		# Handle case where no SVG code could be extracted
		if not rst['actions']:
		env.reward = env.config.format_penalty
		env.total_reward += env.reward
		@@ -427,48 +482,54 @@ class SVGService(BaseService):
		"metrics": metrics,
		"info": info,
		"valid": False,
		"done": True
		"done": False # Changed from True to False to allow training to continue
		}, None

		# Process valid SVG
		env.reward = env.config.format_reward
		# Process SVG (valid or invalid)
		env.reward = env.config.format_reward if svg_is_valid else env.config.format_penalty
		env.total_reward += env.reward
		env.gen_svg_code = rst['actions'][0]
		env.valid_actions = rst['actions']

		try:
		# Process SVG to image but don't calculate score yet
		# 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')

		return env_id, {
		"env": env,
		"gen_image": gen_image,
		"gen_image": env.gen_image,
		"gen_svg_code": env.gen_svg_code,
		"rst": rst,
		"metrics": metrics,
		"valid": True,
		"valid": True, # Consider it valid for processing even if SVG is invalid
		"done": False
		}, None

		except Exception as e:
		logging.error(f"Error processing SVG for {env_id}: {e}")
		env.valid_actions = []
		metrics["turn_metrics"]["action_is_valid"] = False
		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": None,
		"gen_svg_code": env.gen_svg_code,
		"rst": rst,
		"metrics": metrics,
		"info": info,
		"valid": False,
		"done": True
		"done": False # Changed from True to False to allow training to continue
		}, None

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

		# Submit all action processing tasks
		@@ -483,7 +544,8 @@ class SVGService(BaseService):
		env_id, result, error = future.result()
		if error:
		logging.error(f"Error processing action for environment {env_id}: {error}")
		error_results[env_id] = ({}, 0.0, True, {"error": 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

+4 −4

Original line number	Diff line number	Diff line
		@@ -18,7 +18,6 @@ def is_valid_svg(svg_text):
		svgstr2paths(svg_text)
		return True
		except Exception as e:
		print(f"Invalid SVG: {str(e)}")
		return False

		def clean_svg(svg_text, output_width=None, output_height=None):
		@@ -59,15 +58,16 @@ def use_placeholder():

		def process_and_rasterize_svg(svg_string, resolution=256, dpi=128, scale=2):
		try:
		svgstr2paths(svg_string) # This will raise an exception if the svg is still not valid
		svgstr2paths(svg_string)
		out_svg = svg_string
		except:
		try:
		svg = clean_svg(svg_string)
		svgstr2paths(svg) # This will raise an exception if the svg is still not valid
		svgstr2paths(svg) # Try again with cleaned SVG
		out_svg = svg
		except Exception as e:
		out_svg = use_placeholder()
		print(f"SVG processing failed: {e}")
		out_svg = """<svg width="{0}" height="{0}" xmlns="http://www.w3.org/2000/svg"><rect width="100%" height="100%" fill="white"/></svg>""".format(resolution)

		raster_image = rasterize_svg(out_svg, resolution, dpi, scale)
		return out_svg, raster_image

vagen/env/svgdino/svg_utils.py

+4 −3

Original line number	Diff line number	Diff line
		@@ -59,15 +59,16 @@ def use_placeholder():

		def process_and_rasterize_svg(svg_string, resolution=256, dpi=128, scale=2):
		try:
		svgstr2paths(svg_string) # This will raise an exception if the svg is still not valid
		svgstr2paths(svg_string)
		out_svg = svg_string
		except:
		try:
		svg = clean_svg(svg_string)
		svgstr2paths(svg) # This will raise an exception if the svg is still not valid
		svgstr2paths(svg) # Try again with cleaned SVG
		out_svg = svg
		except Exception as e:
		out_svg = use_placeholder()
		print(f"SVG processing failed: {e}")
		out_svg = """<svg width="{0}" height="{0}" xmlns="http://www.w3.org/2000/svg"><rect width="100%" height="100%" fill="white"/></svg>""".format(resolution)

		raster_image = rasterize_svg(out_svg, resolution, dpi, scale)
		return out_svg, raster_image

Admin message