Commit e6a619f2 authored by jameskrw's avatar jameskrw Committed by YaningGao
Browse files

updated crossview qa

parent 01e9c53e
Loading
Loading
Loading
Loading
+9 −0
Original line number Diff line number Diff line
@@ -74,6 +74,15 @@ python vagen/env/alfworld/startx.py 0
python vagen/server/server.py
```



### crossview
cd vagen/env/crossview
git clone https://huggingface.co/datasets/yinbq/CrossViewQA
cd CrossViewQA
mkdir -p extracted_images
unzip other_all_image.zip -d extracted_images

## Benchmark your Env and Service
env/service running time varies on different devices, you can benchmark current env/service or debug your own env/service as follow:
### Start a env benchmark
+194 −0
Original line number Diff line number Diff line
from vagen.env.base.base_env import BaseEnv
from vagen.env.base.base_env_config import BaseEnvConfig
from typing import Dict, List, Tuple, Any, Optional
import json
import os
import random
import re
from PIL import Image
from dataclasses import dataclass, field
from .env_config import CrossViewQAEnvConfig
from vagen.env.utils.context_utils import parse_llm_raw_response


class CrossViewQAEnv(BaseEnv):
    def __init__(self, config: CrossViewQAEnvConfig):
        self.config = config
        self.script_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)),"CrossViewQA")
        self.data_path = os.path.join(self.script_dir, config.data_path)
        self.image_dir = os.path.join(self.script_dir, config.image_dir)
        
        # Load dataset
        with open(self.data_path, 'r', encoding='utf-8') as f:
            self.dataset = json.load(f)
        print(f"Loaded {len(self.dataset)} examples from {self.data_path}")
        
        self.current_data = None
        self.current_seed = None
        self.done = False
        self.total_reward = 0
    
    def reset(self, seed=None) -> Tuple[Dict, Dict]:
        """Reset environment with new seed"""
        if seed is not None:
            self.current_seed = seed
            random.seed(seed)
        
        self.done = False
        self.total_reward = 0
        
        # Select a random data point
        idx = self.current_seed % len(self.dataset) if self.current_seed is not None else random.randint(0, len(self.dataset) - 1)
        self.current_data = self.dataset[idx]
        
        # Create observation
        obs = self._create_observation()
        info = {
            "ground_truth": self.current_data["conversation"][1]["content"],
            "question_id": self.current_data["id"],
        }
        
        return obs,info
    
    def _create_observation(self) -> Dict:
        """Create observation with question and images"""
        # Get question from conversation
        question = self.current_data["conversation"][0]["content"]
        
        # Load images
        images = []
        for path in self.current_data["images"]:
            # Handle path that starts with other_all_image/
            
            full_path = os.path.join(self.image_dir, path)
            img = Image.open(full_path)
           
            img = img.resize(self.config.image_size, Image.LANCZOS)
            images.append(img)
           
        
        # Create observation string with image placeholders
        image_placeholders = " ".join([self.config.image_placeholder] * len(images))
        obs_str = f"Question: {question}\n{image_placeholders}\nPlease look at the images and answer the question."
        
        return {
            'obs_str': obs_str,
            'multi_modal_data': {
                self.config.image_placeholder: images
            }
        }
    
    def step(self, llm_raw_response) -> Tuple[Dict, float, bool, Dict]:
        """Process the LLM's response and compute reward"""
 
        # Parse the response
        parsed_response = parse_llm_raw_response(
            llm_raw_response,
            special_token_list=self.config.special_token_list,
            action_sep=self.config.action_sep
        )
        
        # Get action content and ground truth
        action_content = parsed_response["action_content"].strip()
        ground_truth = self.current_data["conversation"][1]["content"].strip()
        
        # Simple exact match (case-insensitive)
        action_is_valid = action_content != ""
        success = action_is_valid and action_content.lower() == ground_truth.lower()
        action_is_effective = action_is_valid
        
        # Compute reward - base reward + format reward if applicable
        reward = 5.0 if success else 0.0
        if parsed_response["format_correct"] and action_is_valid:
            reward += self.config.format_reward
        
        self.total_reward += reward
        
        # Set done to True (single-step environment)
        self.done = True
        
        # Return observation, reward, done, info
        obs = self._create_observation()
        
        info = {
            "metrics":{ 
                "turn_metrics": {
                "action_is_effective": action_is_effective,
                "action_is_valid": action_is_valid,
            },
                "traj_metrics": {
                    "success": success,  # Will be set to True if agent reaches goal
                }
            },
            "llm_raw_response": llm_raw_response,
            "llm_response": parsed_response["llm_response"],
            "think_content": parsed_response["think_content"],
            "action_content": action_content,
            "actions": parsed_response["actions"],
            "ground_truth": ground_truth,
        }
        
        return obs, reward, self.done, info
    
    def close(self):
        """Close the environment"""
        pass
    
    def system_prompt(self) -> str:
        """Get the system prompt for the environment"""
        return "You are an AI assistant that answers visual questions based on images. " \
               "Given images and a question, first think through the problem in the <think> section, " \
               "and then provide your final answer in the <answer> section."
    
    def compute_reward(self) -> float:
        """Return the total reward accumulated so far"""
        return self.total_reward


if __name__ == "__main__":
    # Create config
    config = CrossViewQAEnvConfig()
    
    # Create environment
    env = CrossViewQAEnv(config)
    
    print("System prompt:")
    print(env.system_prompt())
    print("\n" + "-"*50 + "\n")
    
    
    i = 0
    while True:
        # Get user input
        # Reset environment and get first observation
        obs, info = env.reset(seed=i)
        print("Question:")
        print(obs["obs_str"])
        print("\nGround truth:", info["ground_truth"])
        if config.image_placeholder in obs["multi_modal_data"] and obs["multi_modal_data"][config.image_placeholder]:
            os.makedirs("./test_crossview", exist_ok=True)
            for j, img in enumerate(obs["multi_modal_data"][config.image_placeholder]):
                img.save(f"./test_crossview/crossview_{i}_{j}.png")
        print(f"\nSaved {len(obs['multi_modal_data'][config.image_placeholder])} images to ./test_crossview/")
        answer = input("\nEnter your answer (or just press Enter to use the default format): ")
        
        # If user just pressed Enter, use a default think/answer format
        if not answer:
            llm_response = "<think>Analyzing the two views...</think><answer>B</answer>"
        # If answer doesn't have the think/answer format, add it
        elif "<think>" not in answer:
            llm_response = f"<think>Analyzing the two views...</think><answer>{answer}</answer>"
        else:
            llm_response = answer
        
        # Step the environment
        obs, reward, done, info = env.step(llm_response)
        
        # Display results
        print("\nAction Result:")
        print(f"info: {info}")
        print(f"Reward: {reward}")
        print(f"Total Reward: {env.compute_reward()}")
        i+=1
    
  
 No newline at end of file
+11 −0
Original line number Diff line number Diff line
from vagen.env.base.base_env_config import BaseEnvConfig
from dataclasses import dataclass, fields,field
from typing import Optional, List, Union
@dataclass
class CrossViewQAEnvConfig(BaseEnvConfig):
    data_path: str = "crossviewQA_train_qwenformat_singleletter.json"
    image_dir: str = "extracted_images"
    image_size: tuple = (300, 300)
    
    def config_id(self) -> str:
        return f"CrossViewQAEnv"
 No newline at end of file
+1 −3
Original line number Diff line number Diff line
@@ -180,12 +180,10 @@ class FrozenLakeEnv(BaseEnv):
        if metrics["turn_metrics"]['action_is_valid'] and rst["format_correct"]:
            self.reward += self.config.format_reward
        
        # Add metrics to info dictionary
        info["metrics"] = metrics
        
        # Check if position changed to determine if action was effective
        metrics["turn_metrics"]['action_is_effective'] = not np.array_equal(prev_player_position, self._get_player_position())
        
        info["metrics"] = metrics
        # Update total reward for the episode
        self.total_reward += self.reward
        
+3 −0
Original line number Diff line number Diff line
@@ -18,6 +18,8 @@ def parse_llm_raw_response(response: str,special_token_list=None,action_sep=',',

    pattern = r'<think>(.*?)</think>\s*<answer>(.*?)</answer>'
    match = re.search(pattern, response, re.DOTALL)
    format_correct = match is not None
    
    if not match:
        think_content, action_content, actions = "", "", []
    else:
@@ -38,6 +40,7 @@ def parse_llm_raw_response(response: str,special_token_list=None,action_sep=',',
        "think_content": think_content,
        "action_content": action_content,
        "actions": actions,
        "format_correct": format_correct,
    }