batch reward model and minor (63e305e0) · Commits · 张泽凯 / VAGEN

vagen/env/server.py

+0 −79

Original line number	Diff line number	Diff line
		@@ -62,13 +62,9 @@ class BatchEnvServer:
		return jsonify({"error": "Missing required parameter: ids2configs"}), 400

		ids2configs = data['ids2configs']
		try:
		self._create_environments_batch(ids2configs)
		return jsonify({"success": True}), 200
		except Exception as e:
		return jsonify({"error": str(e)}), 500

		# Batch endpoints aligned with service interface
		@self.app.route('/batch/reset', methods=['POST'])
		def reset_batch():
		"""Reset multiple environments endpoint"""
		@@ -77,11 +73,8 @@ class BatchEnvServer:
		return jsonify({"error": "Missing required parameter: ids2seeds"}), 400

		ids2seeds = data['ids2seeds']
		try:
		results = self._reset_batch(ids2seeds)
		return jsonify({"results": results}), 200
		except Exception as e:
		return jsonify({"error": str(e)}), 500

		@self.app.route('/batch/step', methods=['POST'])
		def step_batch():
		@@ -91,11 +84,8 @@ class BatchEnvServer:
		return jsonify({"error": "Missing required parameter: ids2actions"}), 400

		ids2actions = data['ids2actions']
		try:
		results = self._step_batch(ids2actions)
		return jsonify({"results": results}), 200
		except Exception as e:
		return jsonify({"error": str(e)}), 500

		@self.app.route('/batch/reward', methods=['POST'])
		def compute_reward_batch():
		@@ -105,11 +95,8 @@ class BatchEnvServer:
		return jsonify({"error": "Missing required parameter: env_ids"}), 400

		env_ids = data['env_ids']
		try:
		rewards = self._compute_reward_batch(env_ids)
		return jsonify({"rewards": rewards}), 200
		except Exception as e:
		return jsonify({"error": str(e)}), 500

		@self.app.route('/batch/system_prompt', methods=['POST'])
		def get_system_prompts_batch():
		@@ -119,11 +106,8 @@ class BatchEnvServer:
		return jsonify({"error": "Missing required parameter: env_ids"}), 400

		env_ids = data['env_ids']
		try:
		prompts = self._get_system_prompts_batch(env_ids)
		return jsonify({"system_prompts": prompts}), 200
		except Exception as e:
		return jsonify({"error": str(e)}), 500

		@self.app.route('/batch/close', methods=['POST'])
		def close_batch():
		@@ -133,28 +117,20 @@ class BatchEnvServer:
		return jsonify({"error": "Missing required parameter: env_ids"}), 400

		env_ids = data['env_ids']
		try:
		self._close_batch(env_ids)
		return jsonify({"status": "success"}), 200
		except Exception as e:
		return jsonify({"error": str(e)}), 500

		# Individual environment endpoints (for backward compatibility)
		@self.app.route('/reset/<env_id>', methods=['POST'])
		def reset_environment(env_id):
		"""Reset single environment endpoint"""
		data = request.json or {}
		seed = data.get('seed')

		try:
		results = self._reset_batch({env_id: seed})
		if env_id not in results:
		return jsonify({"error": f"Environment {env_id} not found"}), 404

		obs, info = results[env_id]
		return jsonify({"observation": obs, "info": info}), 200
		except Exception as e:
		return jsonify({"error": str(e)}), 500

		@self.app.route('/step/<env_id>', methods=['POST'])
		def step_environment(env_id):
		@@ -164,8 +140,6 @@ class BatchEnvServer:
		return jsonify({"error": "Missing required parameter: action"}), 400

		action = data['action']

		try:
		results = self._step_batch({env_id: action})
		if env_id not in results:
		return jsonify({"error": f"Environment {env_id} not found"}), 404
		@@ -177,41 +151,28 @@ class BatchEnvServer:
		"done": done,
		"info": info
		}), 200
		except Exception as e:
		return jsonify({"error": str(e)}), 500

		@self.app.route('/reward/<env_id>', methods=['GET'])
		def compute_reward(env_id):
		"""Compute reward for single environment endpoint"""
		try:
		rewards = self._compute_reward_batch([env_id])
		if env_id not in rewards:
		return jsonify({"error": f"Environment {env_id} not found"}), 404

		return jsonify({"reward": rewards[env_id]}), 200
		except Exception as e:
		return jsonify({"error": str(e)}), 500

		@self.app.route('/system_prompt/<env_id>', methods=['GET'])
		def get_system_prompt(env_id):
		"""Get system prompt for single environment endpoint"""
		try:
		prompts = self._get_system_prompts_batch([env_id])
		if env_id not in prompts:
		return jsonify({"error": f"Environment {env_id} not found"}), 404

		return jsonify({"system_prompt": prompts[env_id]}), 200
		except Exception as e:
		return jsonify({"error": str(e)}), 500

		@self.app.route('/close/<env_id>', methods=['DELETE'])
		def close_environment(env_id):
		"""Close single environment endpoint"""
		try:
		self._close_batch([env_id])
		return jsonify({"status": "success"}), 200
		except Exception as e:
		return jsonify({"error": str(e)}), 500

		def _get_service_for_env_name(self, env_name: str) -> BaseService:
		"""
		@@ -233,12 +194,7 @@ class BatchEnvServer:
		raise ValueError(f"No service class registered for environment type: {env_name}")

		service_class = REGISTERED_ENV[env_name]["service_cls"]

		try:
		# Create service instance
		self.services[env_name] = service_class()
		except Exception as e:
		raise ValueError(f"Could not create service for environment type {env_name}: {str(e)}")

		return self.services[env_name]

		@@ -288,14 +244,12 @@ class BatchEnvServer:
		service_to_configs[env_name] = {}
		service_to_configs[env_name][env_id] = config


		# Call create_environments_batch method on each service
		for env_name, configs in service_to_configs.items():
		service = self.services[env_name]
		service.create_environments_batch(configs)



		def _reset_batch(self, ids2seeds: Dict[str, Any]) -> Dict[str, Tuple[Any, Any]]:
		"""
		Reset multiple environments.
		@@ -309,14 +263,10 @@ class BatchEnvServer:
		# Group environment IDs by service
		service_groups = {}
		for env_id, seed in ids2seeds.items():
		try:
		service, env_name = self._get_service_for_env(env_id)
		if env_name not in service_groups:
		service_groups[env_name] = (service, {})
		service_groups[env_name][1][env_id] = seed
		except ValueError as e:
		print(f"Environment ID '{env_id}' not found: {e}")
		pass

		# Reset environments through respective services
		results = {}
		@@ -339,14 +289,10 @@ class BatchEnvServer:
		# Group environment IDs by service
		service_groups = {}
		for env_id, action in ids2actions.items():
		try:
		service, env_name = self._get_service_for_env(env_id)
		if env_name not in service_groups:
		service_groups[env_name] = (service, {})
		service_groups[env_name][1][env_id] = action
		except ValueError as e:
		print(f"Environment ID '{env_id}' not found: {e}")
		pass

		# Step environments through respective services
		results = {}
		@@ -369,14 +315,10 @@ class BatchEnvServer:
		# Group environment IDs by service
		service_groups = {}
		for env_id in env_ids:
		try:
		service, env_name = self._get_service_for_env(env_id)
		if env_name not in service_groups:
		service_groups[env_name] = (service, [])
		service_groups[env_name][1].append(env_id)
		except ValueError as e:
		print(f"Environment ID '{env_id}' not found: {e}")
		pass

		# Compute rewards through respective services
		results = {}
		@@ -399,16 +341,10 @@ class BatchEnvServer:
		# Group environment IDs by service
		service_groups = {}
		for env_id in env_ids:
		try:
		service, env_name = self._get_service_for_env(env_id)
		if env_name not in service_groups:
		service_groups[env_name] = (service, [])
		service_groups[env_name][1].append(env_id)
		except ValueError:
		print(f"Trying to get service for env_id: {env_id}")
		print(f"Available env_ids: {list(self.env_to_service.keys())}")
		# Environment not found, skip it
		pass

		# Get system prompts through respective services
		results = {}
		@@ -428,17 +364,12 @@ class BatchEnvServer:
		# Group environment IDs by service
		service_groups = {}
		for env_id in env_ids:
		try:
		service, env_name = self._get_service_for_env(env_id)
		if env_name not in service_groups:
		service_groups[env_name] = (service, [])
		service_groups[env_name][1].append(env_id)

		# Remove from tracking
		del self.env_to_service[env_id]
		except ValueError as e:
		print(f"Environment ID '{env_id}' not found: {e}")
		pass

		# Close environments through respective services
		for env_name, (service, group_env_ids) in service_groups.items():
		@@ -476,15 +407,11 @@ class BatchEnvServer:
		retry_delay = 0.5
		for _ in range(max_retries):
		time.sleep(retry_delay)
		try:
		import requests
		response = requests.get(f"http://{self.host}:{self.port}/health", timeout=1)
		if response.status_code == 200:
		print(f"Server started on http://{self.host}:{self.port}")
		break
		except Exception as e:
		print(f"Error while checking server health: {e}")
		pass
		else:
		print("Server may not have started properly")
		else:
		@@ -507,14 +434,8 @@ class BatchEnvServer:
		# Shut down the Flask server
		self.is_running = False
		if self.server_thread and self.server_thread.is_alive():
		# This doesn't actually stop Flask in a clean way
		# In a production environment, you would use a proper WSGI server
		import requests
		try:
		requests.post(f"http://{self.host}:{self.port}/shutdown")
		except Exception as e:
		print(f"Error while shutting down server: {e}")
		pass

		print("Server stopped")

vagen/env/svg/config.py

+1 −0

Original line number	Diff line number	Diff line
		@@ -9,6 +9,7 @@ class SVGConfig(BaseConfig):
		data_dir: str = "vagen/env/svg/data"
		seed: int = 16
		split: str = "train"
		action_sep: str = "~~"
		# Score configuration
		model_size: str = "small" # 'small', 'base', or 'large'
		dino_only: bool = False

vagen/env/svg/dino.py

+42 −0

Original line number	Diff line number	Diff line
		@@ -6,6 +6,8 @@ from PIL import Image
		import torch.nn as nn
		import threading
		import logging
		from typing import Dict, List, Tuple, Optional, Any, Union


		# @TODO clean codes of this section

		@@ -124,6 +126,19 @@ class DINOScoreCalculator(BaseMetric):
		return AutoModel.from_pretrained(model_size), AutoImageProcessor.from_pretrained(model_size)

		def process_input(self, image, processor):
		if isinstance(image, list):
		if all(isinstance(img, Image.Image) for img in image):
		with torch.no_grad():
		inputs = processor(images=image, return_tensors="pt").to(self.device)
		outputs = self.model(**inputs)
		features = outputs.last_hidden_state.mean(dim=1)
		return features
		else:
		features_list = []
		for img in image:
		features_list.append(self.process_input(img, processor))
		return torch.cat(features_list, dim=0)

		if isinstance(image, str):
		image = Image.open(image)
		if isinstance(image, Image.Image):
		@@ -148,3 +163,30 @@ class DINOScoreCalculator(BaseMetric):
		sim = (sim + 1) / 2

		return sim

		def calculate_batch_scores(self, gt_images: List[Any], gen_images: List[Any]) -> List[float]:
		"""
		Calculate similarity scores for multiple image pairs in a single batch

		Args:
		gt_images: List of ground truth images (PIL Images, file paths, or tensors)
		gen_images: List of generated images (PIL Images, file paths, or tensors)

		Returns:
		List of similarity scores (float values between 0-1)
		"""
		if not gt_images:
		return []

		gt_features = self.process_input(gt_images, self.processor)

		gen_features = self.process_input(gen_images, self.processor)

		cos = nn.CosineSimilarity(dim=1)
		similarities = cos(gt_features, gen_features)

		scores = [(sim.item() + 1) / 2 for sim in similarities]

		return scores


		No newline at end of file

vagen/env/svg/score.py

+89 −0

Original line number	Diff line number	Diff line
		@@ -85,3 +85,92 @@ def calculate_total_score(gt_im, gen_im, gt_code, gen_code, score_config, dino_m
		scores['total_score'] = max(0.0, weighted_sum)

		return scores

		def calculate_total_score_batch(gt_images, gen_images, gt_codes, gen_codes, score_configs, dino_model=None):
		"""
		Calculate scores for multiple image pairs in batch mode

		Args:
		gt_images: List of ground truth images
		gen_images: List of generated images
		gt_codes: List of ground truth SVG codes
		gen_codes: List of generated SVG codes
		score_configs: List of scoring parameters dictionaries
		dino_model: Pre-loaded DINO model (optional)

		Returns:
		List of dictionaries containing all scores
		"""
		batch_size = len(gt_images)
		if batch_size == 0:
		return []

		# Verify all inputs have same batch size
		if not (len(gen_images) == len(gt_codes) == len(gen_codes) == len(score_configs) == batch_size):
		raise ValueError("All input lists must have the same length")

		# Initialize results
		batch_results = [{
		"dino_score": 0.0,
		"structural_score": 0.0,
		"total_score": 0.0
		} for _ in range(batch_size)]

		# Check if we need to calculate DINO scores
		need_dino = any(score_config.get("dino_weight", 5.0) > 0 for score_config in score_configs)

		# Calculate DINO scores in batch if needed
		if need_dino:
		if dino_model is None:
		from vagen.env.svg.dino import get_dino_model
		# Default to small model size if not specified
		model_size = score_configs[0].get("model_size", "small") if score_configs else "small"
		dino_model = get_dino_model(model_size)

		# Calculate all DINO scores at once using batch processing
		dino_scores = dino_model.calculate_batch_scores(gt_images, gen_images)

		# Assign scores to results
		for i, score in enumerate(dino_scores):
		batch_results[i]["dino_score"] = float(score)

		# Calculate structural scores and total scores
		for i in range(batch_size):
		score_config = score_configs[i]
		result = batch_results[i]

		# Check if DINO-only mode
		dino_only = score_config.get("dino_only", False)
		if dino_only:
		result["total_score"] = result["dino_score"]
		continue

		# Get model size for default weights
		model_size = score_config.get("model_size", "small")

		# Define default weights based on model size
		default_weights = {
		"small": {"dino": 3.0, "structural": 7.0},
		"base": {"dino": 5.0, "structural": 5.0},
		"large": {"dino": 6.0, "structural": 4.0}
		}

		# Get weights with defaults
		weights = {
		"dino": score_config.get("dino_weight", default_weights[model_size]["dino"]),
		"structural": score_config.get("structural_weight", default_weights[model_size]["structural"])
		}

		# Calculate structural score if needed
		if weights["structural"] > 0:
		from vagen.env.svg.score import calculate_structural_accuracy
		result["structural_score"] = max(0.0, float(calculate_structural_accuracy(gt_images[i], gen_images[i])))

		# Calculate weighted total score
		weighted_sum = (
		result["dino_score"] * weights["dino"] +
		result["structural_score"] * weights["structural"]
		)
		result["total_score"] = max(0.0, weighted_sum)

		return batch_results
		No newline at end of file

vagen/env/svg/service.py

+185 −42

File changed.

Preview size limit exceeded, changes collapsed.

Admin message