revise inference rollout (3e9ac266) · Commits · 张泽凯 / VAGEN

vagen/mllm_agent/inference_rollout/inference_rollout_service.py

+11 −123

Original line number	Diff line number	Diff line
		# vagen/mllm_agent/inference_rollout/inference_rollout_service.py

		import os
		import time
		from typing import List, Dict, Tuple, Optional, Any
		@@ -166,7 +168,7 @@ class InferenceRolloutService(BaseRollout):
		print(f"All environments completed after {step} steps")
		break

		# Collect prompts and images for active environments
		# Collect prompts for active environments
		env_messages = {}

		for env_id in active_envs:
		@@ -233,6 +235,10 @@ class InferenceRolloutService(BaseRollout):
		"""
		Generate responses for multiple environments.

		This method now properly aligns with the model interface that expects
		message lists directly, and handles multimodal data according to the
		training rollout format.

		Args:
		env_messages: Dictionary mapping environment IDs to conversation histories

		@@ -243,21 +249,13 @@ class InferenceRolloutService(BaseRollout):
		env_ids = list(env_messages.keys())
		prompts = []

		# Collect all messages for batch generation
		for env_id in env_ids:
		messages = env_messages[env_id]
		# Check if any message has multimodal data
		has_images = any("multi_modal_data" in msg for msg in messages)

		if has_images:
		# Process multimodal input
		prompt = self._process_multimodal_messages(messages)
		else:
		# Process text-only input
		prompt = self.model_interface.format_prompt(messages)

		prompts.append(prompt)
		prompts.append(messages) # Model interface expects message lists directly

		# Generate responses for all prompts
		# The model interface will handle multimodal data extraction internally
		batch_results = self.model_interface.generate(prompts)

		# Extract responses
		@@ -266,39 +264,6 @@ class InferenceRolloutService(BaseRollout):

		return responses

		def _process_multimodal_messages(self, messages: List[Dict]) -> Dict:
		"""
		Process messages with multimodal data for the model interface.

		Args:
		messages: List of message dictionaries

		Returns:
		Processed input for the model interface
		"""
		# Extract images from messages
		all_images = []
		for message in messages:
		if "multi_modal_data" in message:
		for key, values in message["multi_modal_data"].items():
		for value in values:
		# Handle different image formats
		if isinstance(value, PIL.Image.Image):
		all_images.append(value)
		# Handle serialized images from the service
		elif isinstance(value, dict) and "__pil_image__" in value:
		from vagen.server.serial import deserialize_pil_image
		all_images.append(deserialize_pil_image(value))

		# Process images
		processed_images = self.model_interface.process_images(all_images)

		# Return formatted input
		return {
		"messages": messages,
		"images": processed_images
		}

		def recording_to_log(self) -> List[Dict]:
		"""
		Format and return results in a format compatible with logging.
		@@ -348,6 +313,7 @@ class InferenceRolloutService(BaseRollout):
		for message in self.recordings[env_id]:
		if "multi_modal_data" in message:
		for key, values in message["multi_modal_data"].items():
		if key == "<image>" or "image" in key.lower():
		for value in values:
		# Handle different image formats
		if isinstance(value, PIL.Image.Image):
		@@ -421,81 +387,3 @@ class InferenceRolloutService(BaseRollout):

		if self.debug:
		print("Closed all environments and cleaned up resources")
		No newline at end of file


		def run_test():
		"""Run a simple test of the InferenceRolloutService."""
		print("=== Testing InferenceRolloutService ===")

		# Create a model interface
		model = MockModelInterface()

		# Create a configuration
		config = {
		"max_steps": 5,
		"show_progress": True,
		"debug": True
		}

		# Create the service
		service = InferenceRolloutService(
		config=config,
		model_interface=model,
		debug=True
		)

		# Create environment configurations
		env_configs = [
		{
		"env_name": "frozenlake",
		"env_config": {
		"size": 4,
		"is_slippery": False,
		"render_mode": "vision"
		},
		"seed": 42
		},
		{
		"env_name": "frozenlake",
		"env_config": {
		"size": 4,
		"is_slippery": True,
		"render_mode": "text"
		},
		"seed": 43
		}
		]

		try:
		# Reset environments
		print("\nResetting environments...")
		service.reset(env_configs)

		# Run inference
		print("\nRunning inference...")
		service.run()

		# Get results
		print("\nGetting results...")
		results = service.recording_to_log()

		# Print results
		print("\n=== Results ===")
		for result in results:
		print(f"Environment: {result['env_id']}")
		print(f"Config ID: {result['config_id']}")
		print(f"Steps: {result['metrics']['step']}")
		print(f"Done: {result['metrics']['done']}")
		print(f"Score: {result['metrics']['score']}")
		print("Metrics:", {k: v for k, v in result['metrics'].items() if k not in ['step', 'done', 'score']})
		print("---")

		finally:
		# Clean up
		print("\nCleaning up...")
		service.close()

		print("=== Test completed ===")

		if __name__ == "__main__":
		run_test()
		No newline at end of file

vagen/mllm_agent/model_interface/init.py

0 → 100644

+8 −0

Original line number	Diff line number	Diff line
		from .vllm import VLLMModelInterface, VLLMModelConfig

		REGISTERED_MODEL = {
		"vllm": {
		"model_cls": VLLMModelInterface,
		"config_cls": VLLMModelConfig,
		},
		}
		No newline at end of file

vagen/mllm_agent/model_interface/base_model_config.py

0 → 100644

+31 −0

Original line number	Diff line number	Diff line
		from abc import ABC, abstractmethod
		from dataclasses import dataclass, field, fields
		from typing import Optional, Dict, Any

		@dataclass
		class BaseModelConfig(ABC):
		"""Abstract base configuration for all model interfaces."""

		# Common parameters across all models
		model_name: str
		max_tokens: int = 1024
		temperature: float = 0.7
		seed: Optional[int] = None

		@abstractmethod
		def config_id(self) -> str:
		"""Generate a unique identifier for this configuration."""
		pass

		def to_dict(self) -> Dict[str, Any]:
		"""Convert config to dictionary for model initialization."""
		return {f.name: getattr(self, f.name) for f in fields(self)}

		def get(self, key, default=None):
		"""
		Get the value of a config key.
		Args:
		key: Key to get
		default: Default value if key is not found
		"""
		return getattr(self, key, default)
		No newline at end of file

vagen/mllm_agent/model_interface/factory_model.py

+115 −56

Original line number	Diff line number	Diff line
		# vagen/mllm_agent/model_interface/factory_model.py

		import logging
		from typing import Dict, Any, Optional

		from .base_model import BaseModelInterface
		from .providers.vllm_model import VLLMModelInterface
		from .base_model_config import BaseModelConfig
		from . import REGISTERED_MODEL

		logger = logging.getLogger(__name__)

		def create_model_interface(config: Dict[str, Any]) -> Optional[BaseModelInterface]:
		class ModelFactory:
		"""
		Factory function to create an appropriate model interface based on configuration.
		Factory for creating and managing model interfaces.
		Uses REGISTERED_MODEL to find and create model instances.
		"""

		@staticmethod
		def create(config: Dict[str, Any]) -> BaseModelInterface:
		"""
		Create a model interface based on configuration.

		Args:
		config: Configuration dictionary containing:
		- provider: Model provider type (e.g., "vllm", "api")
		- model_name: Name of the model to use
		- Other provider-specific parameters
		config: Configuration dictionary with model parameters

		Returns:
		Initialized model interface or None if initialization fails
		Initialized model interface

		Raises:
		ValueError: If provider is unknown or initialization fails
		"""
		provider = config.get("provider", "vllm").lower()
		model_name = config.get("model_name", "")

		logger.info(f"Creating model interface for provider '{provider}' with model '{model_name}'")

		if provider not in REGISTERED_MODEL:
		available_providers = list(REGISTERED_MODEL.keys())
		raise ValueError(f"Unknown provider '{provider}'. Available providers: {available_providers}")

		try:
		if provider == "vllm":
		return VLLMModelInterface(config)
		# Add more providers as needed
		# elif provider == "openai":
		# return OpenAIModelInterface(config)
		# elif provider == "claude":
		# return ClaudeModelInterface(config)
		else:
		logger.error(f"Unknown model provider: {provider}")
		return None
		# Get model and config classes from registry
		model_cls = REGISTERED_MODEL[provider]["model_cls"]
		config_cls = REGISTERED_MODEL[provider]["config_cls"]

		# Create config instance
		model_config = config_cls(**config)

		# Create and return model instance
		return model_cls(model_config)

		except Exception as e:
		logger.error(f"Failed to initialize model interface: {str(e)}")
		return None

		class ModelFactory:
		"""
		Class-based factory for creating and managing model interfaces.
		Provides additional functionality beyond the simple factory function.
		"""
		raise

		@staticmethod
		def create(config: Dict[str, Any]) -> Optional[BaseModelInterface]:
		def create_from_config_instance(config: BaseModelConfig) -> BaseModelInterface:
		"""
		Create a model interface based on configuration.
		Create a model interface from a config instance.

		Args:
		config: Configuration dictionary with model parameters
		config: Model configuration instance

		Returns:
		Initialized model interface or None if initialization fails
		Initialized model interface
		"""
		return create_model_interface(config)
		# Find provider by checking config class type
		for provider, info in REGISTERED_MODEL.items():
		if isinstance(config, info["config_cls"]):
		model_cls = info["model_cls"]
		return model_cls(config)

		raise ValueError(f"No registered provider found for config type {type(config).__name__}")

		@staticmethod
		def get_available_providers() -> Dict[str, Dict[str, Any]]:
		@@ -66,22 +80,27 @@ class ModelFactory:
		Returns:
		Dictionary mapping provider names to their capabilities
		"""
		return {
		"vllm": {
		"description": "Local model inference using vLLM",
		"supports_multimodal": True,
		"supported_models": [
		"Qwen/Qwen2.5-0.5B-Instruct",
		"Qwen/Qwen2.5-VL-3B-Instruct"
		]
		}
		# Add more providers as they become available
		providers = {}

		for provider, info in REGISTERED_MODEL.items():
		provider_info = {
		"model_class": info["model_cls"].__name__,
		"config_class": info["config_cls"].__name__,
		}

		# Add provider-specific info if available
		config_cls = info["config_cls"]
		if hasattr(config_cls, 'get_provider_info'):
		provider_info.update(config_cls.get_provider_info())

		providers[provider] = provider_info

		return providers

		@staticmethod
		def validate_config(config: Dict[str, Any]) -> Dict[str, Any]:
		"""
		Validate and complete a model configuration with defaults if needed.
		Validate and complete a model configuration with defaults.

		Args:
		config: Model configuration to validate
		@@ -89,23 +108,63 @@ class ModelFactory:
		Returns:
		Validated configuration with defaults applied
		"""
		# Copy to avoid modifying the original
		validated = config.copy()
		provider = config.get("provider", "vllm").lower()

		if provider not in REGISTERED_MODEL:
		logger.warning(f"Unknown provider '{provider}' during validation")
		return config

		try:
		# Get config class from registry
		config_cls = REGISTERED_MODEL[provider]["config_cls"]

		# Create config instance (which applies defaults)
		model_config = config_cls(**config)

		# Convert back to dictionary
		return model_config.to_dict()

		except Exception as e:
		logger.warning(f"Could not validate config for provider {provider}: {e}")
		# Return original config if validation fails
		return config

		@staticmethod
		def create_from_yaml_config(yaml_config: Dict[str, Any],
		model_name: str) -> BaseModelInterface:
		"""
		Create a model interface from YAML config format.

		# Set defaults
		if "provider" not in validated:
		validated["provider"] = "vllm"
		Args:
		yaml_config: Full YAML config with potentially multiple models
		model_name: Name of the specific model to create

		Returns:
		Initialized model interface

		Raises:
		KeyError: If model_name not found in config
		"""
		if "models" in yaml_config:
		models_config = yaml_config["models"]
		else:
		models_config = yaml_config

		# Provider-specific validation
		if validated["provider"] == "vllm":
		if "model_name" not in validated:
		validated["model_name"] = "Qwen/Qwen2.5-0.5B-Instruct"
		if model_name not in models_config:
		raise KeyError(f"Model '{model_name}' not found in configuration")

		if "tensor_parallel_size" not in validated:
		validated["tensor_parallel_size"] = 1
		model_config = models_config[model_name]
		return ModelFactory.create(model_config)

		# Ensure Qwen models have trust_remote_code=True
		if "Qwen" in validated.get("model_name", ""):
		validated["trust_remote_code"] = True
		@staticmethod
		def is_provider_supported(provider: str) -> bool:
		"""
		Check if a provider is supported.

		Args:
		provider: Provider name to check

		return validated
		No newline at end of file
		Returns:
		True if provider is supported, False otherwise
		"""
		return provider.lower() in REGISTERED_MODEL
		No newline at end of file

vagen/mllm_agent/model_interface/providers/test_vllm.py

deleted100644 → 0

+0 −124

Original line number	Diff line number	Diff line
		# test_vllm_model.py
		import sys
		import os
		from PIL import Image
		import numpy as np
		import logging
		import json

		# Setup logging
		logging.basicConfig(level=logging.INFO)

		# Add parent directory to path to import the module
		sys.path.append("../")

		from vagen.mllm_agent.model_interface.providers.vllm_model import VLLMModelInterface

		def create_test_image():
		"""Create a simple test image"""
		# Create a 100x100 RGB image with a gradient
		img_array = np.zeros((100, 100, 3), dtype=np.uint8)

		# Create a simple gradient
		for i in range(100):
		for j in range(100):
		img_array[i, j, 0] = i * 255 // 100 # Red channel
		img_array[i, j, 1] = j * 255 // 100 # Green channel
		img_array[i, j, 2] = 100 # Blue channel

		# Convert to PIL Image
		img = Image.fromarray(img_array)
		return img

		def test_text_model():
		"""Test the text-only Qwen model"""
		print("\n===== Testing Qwen2.5-0.5B-Instruct (Text) =====")

		# Configuration for text model
		config = {
		"model_name": "Qwen/Qwen2.5-0.5B-Instruct",
		"model_family": "qwen",
		"max_tokens": 256,
		"temperature": 0.7
		}

		try:
		# Initialize model
		model = VLLMModelInterface(config)
		print(f"Successfully initialized model: {model.get_model_info()['name']}")

		# Test conversation
		messages = [
		{"role": "system", "content": "You are a helpful assistant."},
		{"role": "user", "content": "What is the capital of France?"}
		]

		# Generate response
		print("Generating response...")
		results = model.generate([messages])

		# Print result
		if results and len(results) > 0:
		print("\nResponse:")
		print(results[0]["text"])
		print(f"\nToken usage: {results[0]['usage']}")
		else:
		print("No response generated.")

		except Exception as e:
		print(f"Error testing text model: {str(e)}")

		def test_multimodal_model():
		"""Test the multimodal Qwen model"""
		print("\n===== Testing Qwen2.5-VL-3B-Instruct (Multimodal) =====")

		# Configuration for multimodal model
		config = {
		"model_name": "Qwen/Qwen2.5-VL-3B-Instruct",
		"model_family": "qwen",
		"max_tokens": 256,
		"temperature": 0.7
		}

		try:
		# Initialize model
		model = VLLMModelInterface(config)
		print(f"Successfully initialized model: {model.get_model_info()['name']}")

		# Create test image
		test_image = create_test_image()

		# Test conversation with image
		messages = [
		{"role": "system", "content": "You are a helpful assistant."},
		{"role": "user", "content": "What do you see in this image?"}
		]

		# Generate response with image
		print("Generating response with image...")
		prompt = {
		"messages": messages,
		"images": [test_image]
		}

		results = model.generate([prompt])

		# Print result
		if results and len(results) > 0:
		print("\nResponse:")
		print(results[0]["text"])
		print(f"\nToken usage: {results[0]['usage']}")
		else:
		print("No response generated.")

		except Exception as e:
		print(f"Error testing multimodal model: {str(e)}")

		if __name__ == "__main__":
		# Test text model
		test_text_model()

		# Test multimodal model
		test_multimodal_model()

		print("\nTests completed!")
		No newline at end of file

Admin message