atropos/environments/hack0/ufc_env/ufc_image_env.py

import os
import random
import sys
import traceback
import csv
from typing import List, Optional, Tuple, Any, Dict
import base64
from PIL import Image
import io
from pydantic import Field

from atroposlib.envs.base import BaseEnv, BaseEnvConfig, OpenaiConfig, ScoredDataGroup
from atroposlib.type_definitions import GameHistory, Item
from atroposlib.utils.tokenize_for_trainer import tokenize_for_trainer


class UFCImageEnvConfig(BaseEnvConfig):
    """Configuration for the UFC Image Environment"""
    fighter_stats_path: str = Field(os.path.join(os.path.dirname(__file__), "fighter_stats.csv"), description="Path to fighter stats CSV")
    fight_data_path: str = Field(os.path.join(os.path.dirname(__file__), "large_dataset.csv"), description="Path to large fight dataset CSV")
    image_folder: str = Field(os.path.join(os.path.dirname(__file__), "fighter_images"), description="Path to fighter images folder")
    max_steps: int = Field(1, description="Only one step per fight prediction")
    temperature: float = Field(0.7, description="Temperature for generation diversity")
    top_p: float = Field(0.95, description="Top p for nucleus sampling")


class UFCImageEnv(BaseEnv):
    """UFC Fight Prediction Environment using only fighter images"""
    name = "ufc_image_predictor"
    env_config_cls = UFCImageEnvConfig

    def __init__(self, config: UFCImageEnvConfig, server_configs: List[OpenaiConfig], slurm=True, testing=False):
        super().__init__(config, server_configs, slurm, testing)
        self.fighter_stats = {}
        self.fight_data = []
        self.current_index = 0
        self.inference_server = self.server.servers[0]  # Get first server as inference server

    async def setup(self):
        """Load the fighter stats and fight data"""
        try:
            print("Loading fighter stats from:", self.config.fighter_stats_path)
            with open(self.config.fighter_stats_path, encoding="utf-8") as f:
                reader = csv.DictReader(f)
                self.fighter_stats = {row["name"]: row for row in reader}
            print(f"Loaded stats for {len(self.fighter_stats)} fighters")

            print("Loading fight data from:", self.config.fight_data_path)
            with open(self.config.fight_data_path, encoding="utf-8") as f:
                reader = csv.DictReader(f)
                self.fight_data = list(reader)
            print(f"Loaded {len(self.fight_data)} fights")

            # Filter out fights where either fighter's image is missing
            filtered_fights = []
            missing_images = set()  # Track unique missing images
            for fight in self.fight_data:
                r_fighter = fight["r_fighter"]
                b_fighter = fight["b_fighter"]
                
                # Convert names to image filename format
                r_slug = r_fighter.lower().replace(" ", "-")
                b_slug = b_fighter.lower().replace(" ", "-")
                
                r_image_path = os.path.join(self.config.image_folder, f"{r_slug}.jpg")
                b_image_path = os.path.join(self.config.image_folder, f"{b_slug}.jpg")
                
                if os.path.exists(r_image_path) and os.path.exists(b_image_path):
                    filtered_fights.append(fight)
                else:
                    if not os.path.exists(r_image_path):
                        missing_images.add(r_fighter)
                    if not os.path.exists(b_image_path):
                        missing_images.add(b_fighter)
            
            if missing_images:
                print(f"\nMissing images for {len(missing_images)} fighters. These fights will be skipped.")
            
            self.fight_data = filtered_fights
            print(f"Filtered to {len(self.fight_data)} fights with complete image sets")

        except Exception as e:
            print(f"Error loading data: {e}")
            traceback.print_exc()
            sys.exit(1)

    def get_fighter_image(self, fighter_name):
        """Convert fighter name to image path and return base64 encoded image"""
        try:
            # Convert name to slug format
            slug = fighter_name.lower().replace(" ", "-")
            
            image_path = os.path.join(self.config.image_folder, f"{slug}.jpg")
            if not os.path.exists(image_path):
                return None
                
            # Convert image to base64
            with Image.open(image_path) as img:
                # Convert RGBA to RGB if necessary
                if img.mode == 'RGBA':
                    img = img.convert('RGB')
                buf = io.BytesIO()
                img.save(buf, format="JPEG")
                image_bytes = buf.getvalue()
                return base64.b64encode(image_bytes).decode("utf-8")
        except Exception as e:
            print(f"Error getting image for {fighter_name}: {e}")
            return None

    async def get_next_item(self) -> Optional[Item]:
        """Get the next fight from the dataset"""
        try:
            if self.current_index >= len(self.fight_data):
                return None
            fight = self.fight_data[self.current_index]
            self.current_index += 1
            
            r_fighter = fight["r_fighter"]
            b_fighter = fight["b_fighter"]
            
            # Get base64 encoded images
            r_image = self.get_fighter_image(r_fighter)
            b_image = self.get_fighter_image(b_fighter)
            
            if not r_image or not b_image:
                print(f"Skipping fight {self.current_index} due to missing images")
                return None
            
            # Format the prompt with images
            prompt_text = (
                "🎤 LADIES AND GENTLEMEN! Welcome to the most electrifying show in sports entertainment "
                "Let's break down this matchup that's got everyone talking!\n\n"
                "In the red corner, we have:(YOUR FIRST IMAGE):\n"
                "And in the blue corner: (YOUR SECOND IMAGE):\n\n"
                "Now, act as your favorite fight comentator, I want you to:\n"                
                "create a fight commentary of whats happening in the fight live\n"
                "Give us your best fight commentary! Make it exciting, make it dramatic, make it sound like you're calling the fight live! "
                "Throw in some classic commentator phrases, maybe a 'OH MY GOODNESS!' or two, and definitely some dramatic pauses for effect.\n\n"
                "End your masterpiece with your prediction in this exact format:\n"  
                "[S1]Hello im your host  [S2] And so am i (name) [S1] Wow. Amazing. (laughs) [S2] Lets get started! (coughs)\n\n"
                "The winner should always be annouced with"
                "\\boxed{Red} or \\boxed{Blue}"
                "Or you will receive a score of -1.0"
            )
            
            # Create multimodal prompt with images
            prompt = tuple([
                {
                    "role": "user", 
                    "content": [
                        {"type": "text", "text": prompt_text},
                        {
                            "type": "image_url",
                            "image_url": {"url": f"data:image/jpeg;base64,{r_image}"}
                        },
                        {
                            "type": "image_url",
                            "image_url": {"url": f"data:image/jpeg;base64,{b_image}"}
                        }
                    ]
                }
            ])
            
            winner = fight.get("winner", "")  # Red or Blue
            ground_truth = f"Answer: {winner}" if winner else ""
            
            return (prompt, ground_truth, None)
            
        except Exception as e:
            print(f"Error in get_next_item: {e}")
            traceback.print_exc()
            return None

    async def collect_trajectories(self, item: Item) -> Tuple[List[Tuple[GameHistory, str, Optional[str]]], List[Item]]:
        to_score = []
        to_backlog = []
        
        system_msg = {
            "role": "system",
            "content": (
                "You are an expert MMA analyst. You will be given two UFC fighters' images. "
                "Your task is to predict the winner of the fight based on their appearance and physique.\n\n"
                "IMPORTANT: You MUST format your response in exactly two parts:\n"
                "1. First, analyze the fighters' appearances and create a fight commentary\n"
                "2. Then on a new line, give ONLY your final prediction in this exact format:\n"
                "\\boxed{Red} or \\boxed{Blue}\n\n"
                "For example:\n"
                "After analyzing the fighters' appearances... [your analysis here]\n"
                "\\boxed{Red}\n\n"
                "If you do not end your response with the \\boxed{} format containing either 'Red' or 'Blue', you will receive a score of -1.0."
            )
        }
        
        user_msg = {
            "role": "user",
            "content": dict(item[0][0])["content"]
        }
        
        messages = [system_msg, user_msg]
        
        try:
            chat_completions = await self.inference_server.chat_completion(
                messages=messages,
                n=self.config.group_size,
                max_tokens=2048,
                temperature=self.config.temperature,
                top_p=self.config.top_p,
                timeout=60,
            )
            for choice in chat_completions.choices:
                assistant_msg = {"role": "assistant", "content": choice.message.content}
                history = [
                    {"role": "system", "content": system_msg["content"]},
                    {"role": "user", "content": user_msg["content"]},
                    {"role": "assistant", "content": choice.message.content}
                ]
                to_score.append((history, item[1], None))
        except Exception as e:
            print(f"Error in collect_trajectories: {e}")
            traceback.print_exc()
            to_backlog.append(item)
            
        if not to_score:
            return None, to_backlog
            
        scored_data = await self.score(to_score)
        return scored_data, to_backlog

    async def score(self, rollout_group_data) -> Optional[ScoredDataGroup]:
        if not rollout_group_data:
            return None
            
        scores = ScoredDataGroup()
        scores["tokens"] = []
        scores["masks"] = []
        scores["scores"] = []
        scores["advantages"] = None
        scores["ref_logprobs"] = None
        scores["messages"] = None
        scores["group_overrides"] = {"group_size": self.config.group_size}
        scores["overrides"] = None
        scores["ground_truths"] = []
        
        random.shuffle(rollout_group_data)
        for item in rollout_group_data:
            out = tokenize_for_trainer(self.tokenizer, item[0])
            tokens = out["tokens"]
            masks = out["masks"]
            
            try:
                # Extract prediction and ground truth
                reply = item[0][-1]["content"]
                ground_truth = item[1].strip().lower()
                
                # Extract color from ground truth (format: "answer: color")
                ground_truth_color = ground_truth.replace("answer:", "").strip()
                
                # Extract color from \boxed{color} format
                import re
                boxed_match = re.search(r"\\boxed{([^}]+)}", reply)
                if boxed_match:
                    prediction = boxed_match.group(1).strip().lower()
                    # Compare just the colors
                    reward = 1.0 if prediction == ground_truth_color else -1.0
                else:
                    # No boxed answer found
                    reward = -1.0
                    
            except Exception as e:
                print(f"Error scoring response: {e}")
                reward = -1.0
                ground_truth = item[1] if isinstance(item[1], str) else ""
                
            if len([i for i in masks if i != -100]) < 10:
                continue
                
            scores["tokens"].append(tokens)
            scores["masks"].append(masks)
            scores["scores"].append(reward)
            scores["ground_truths"].append(ground_truth)
            
            if len(scores["tokens"]) >= self.config.group_size:
                break
                
        if not scores["tokens"]:
            return None
            
        return scores

    async def evaluate(self, *args, **kwargs):
        """No-op evaluation"""
        return

    @classmethod
    def config_init(cls) -> Tuple[BaseEnvConfig, List[OpenaiConfig]]:
        """Initialize configuration for the environment"""
        if not os.environ.get("OPENAI_API_KEY"):
            print("ERROR: OPENAI_API_KEY environment variable is not set!")
            sys.exit(1)

        config = UFCImageEnvConfig(
            wandb_name="ufc_image",
            tokenizer_name="gpt2",
            group_size=2,
            use_wandb=False,
            max_num_workers=2,
            rollout_server_url="http://localhost:8000",
            total_steps=1000,
            batch_size=1,
            steps_per_eval=10,
            ensure_scores_are_not_same=False,
        )

        server_configs = [
            OpenaiConfig(
                model_name="gpt-4o",
                base_url=None,
                api_key=os.environ.get("OPENAI_API_KEY"),
                num_requests_for_eval=1,
            ),
        ]

        return config, server_configs

if __name__ == "__main__":
    UFCImageEnv.cli()