Upload stratego\main.py with huggingface_hub

stratego//main.py

@@ -0,0 +1,431 @@
+import argparse
+import os
+import re
+import time
+import random
+# from stratego.prompt_optimizer import improve_prompt
+from stratego.env.stratego_env import StrategoEnv
+from stratego.prompts import get_prompt_pack
+from stratego.utils.parsing import extract_board_block_lines, extract_legal_moves, extract_forbidden
+from stratego.utils.game_move_tracker import GameMoveTracker as MoveTrackerClass
+from stratego.utils.move_processor import process_move
+from stratego.game_logger import GameLogger
+from stratego.game_analyzer import analyze_and_update_prompt
+from stratego.datasets import auto_push_after_game
+
+
+#Revised to set temperature(13 Nov 2025)
+def build_agent(spec: str, prompt_name: str):
+    """ 
+    Creates and configures an AI agent based on the input string.
+    Example spec: 'ollama:phi3:3.8b'
+    """
+    kind, name = spec.split(":", 1)  # Split string to get model type and name
+    
+    if kind == "ollama":
+        from stratego.models.ollama_model import OllamaAgent
+        # Define the temperature value explicitly
+        AGENT_TEMPERATURE = 0.2
+        
+        # Create the Ollama agent
+        agent = OllamaAgent(
+            model_name=name, 
+            temperature=AGENT_TEMPERATURE, 
+            num_predict=128,  # Allow enough tokens for a complete move response
+            prompt_pack=get_prompt_pack(prompt_name) # Load strategy prompt
+        )
+        
+        # Store temperature for logging
+        agent.temperature = AGENT_TEMPERATURE
+        
+        return agent
+    if kind == "hf":
+        from stratego.models.hf_model import HFLocalAgent
+        return HFLocalAgent(model_id=name, prompt_pack=prompt_name)
+    raise ValueError(f"Unknown agent spec: {spec}")
+
+def print_board(observation: str, size: int = 10):
+    block = extract_board_block_lines(observation, size)
+    if block:
+        print("\n".join(block))
+
+# --- Main Command Line Interface (CLI) ---
+def cli():
+    DEFAULT_ENV = "Stratego-v0"
+    DUEL_ENV = "Stratego-duel"
+    CUSTOM_ENV = "Stratego-custom"
+    tracker = MoveTrackerClass()
+    p = argparse.ArgumentParser()
+    p.add_argument("--p0", default="ollama:deepseek-r1:32b")
+    p.add_argument("--p1", default="ollama:gemma3:1b")
+    # UPDATED HELP TEXT to explain how this parameter relates to VRAM utilization
+    # For large models (120B, 70B), you MUST set this value based on available VRAM(13 Nov 2025)
+    # UPDATED GPU arguments for VRAM control (now defaults to CPU-only)
+    p.add_argument("--p0-num-gpu", type=int, default=0,
+                    help="Number of GPU layers to offload for Player 0. Default is 0 (CPU-only mode). Use a positive number (e.g., 50) to offload layers to GPU/VRAM, or 999 for maximum GPU use.")
+    p.add_argument("--p1-num-gpu", type=int, default=0,
+                    help="Number of GPU layers to offload for Player 1. Default is 0 (CPU-only mode). Use a positive number (e.g., 40) to offload layers to GPU/VRAM, or 999 for maximum GPU use.")    
+    #(13 Nov 2025) NOTE: Default env_id is used as a flag to trigger the interactive menu
+    p.add_argument("--prompt", default="base", help="Prompt preset name (e.g. base|concise|adaptive)")
+    p.add_argument("--env_id", default=DEFAULT_ENV, help="TextArena environment id")
+    p.add_argument("--log-dir", default="logs", help="Directory for per-game CSV logs")
+    p.add_argument("--game-id", default=None, help="Optional custom game id in CSV filename")
+    p.add_argument("--size", type=int, default=10, help="Board size NxN")
+    p.add_argument("--max-turns", type=int, default=None, help="Maximum turns before stopping (for testing). E.g., --max-turns 10")
+
+    args = p.parse_args()
+
+    #(13 Nov 2025) --- INTERACTIVE ENVIRONMENT SELECTION ---
+    if args.env_id == DEFAULT_ENV:
+        print("\n--- Stratego Version Selection ---")
+        print(f"1. Standard Game ({DEFAULT_ENV})")
+        print(f"2. Duel Mode ({DUEL_ENV})")
+        print(f"3. Custom Mode ({CUSTOM_ENV})")
+        
+        while True:
+            choice = input("Enter your choice (1, 2, or 3): ").strip()
+            if not choice or choice == '1':
+                print(f"Selected: {DEFAULT_ENV}")
+                break
+            elif choice == '2':
+                args.env_id = DUEL_ENV
+                args.size = 6
+                print(f"Selected: {DUEL_ENV}")
+                break
+            elif choice == '3':
+                # [CHANGE] Updated prompt range description
+                board = input("Please enter your custom board size in range of 4~9: ").strip()
+                # [CHANGE] Added '4' and '5' to valid options
+                if board in ['4', '5', '6', '7', '8', '9']:
+                    args.env_id = CUSTOM_ENV
+                    args.size = int(board)
+                    print(f"Selected: {CUSTOM_ENV} with size {args.size}x{args.size}")
+                    break
+                else:
+                    print("Invalid choice.")
+            else:
+                print("Invalid choice.")
+
+    # --- Setup Game ---
+    agents = {
+        0: build_agent(args.p0, args.prompt),
+        1: build_agent(args.p1, args.prompt),
+    }
+    # Check if it is really normal Stratego version
+    if (args.env_id == CUSTOM_ENV): 
+        env = StrategoEnv(env_id=CUSTOM_ENV, size=args.size)
+        game_type = "custom"
+    elif (args.env_id == DUEL_ENV):
+        env = StrategoEnv(env_id=DUEL_ENV)
+        game_type = "duel"
+        args.size = 6  # Duel mode uses 6x6 board
+    else:
+        env = StrategoEnv()
+        game_type = "standard"
+    env.reset(num_players=2)
+    
+    # Track game start time
+    game_start_time = time.time()
+    
+    # Simple move history tracker (separate for each player)
+    move_history = {0: [], 1: []}
+
+    with GameLogger(out_dir=args.log_dir, game_id=args.game_id, prompt_name=args.prompt, game_type=game_type, board_size=args.size) as logger:
+        for pid in (0, 1):
+            if hasattr(agents[pid], "logger"):
+                agents[pid].logger = logger
+                agents[pid].player_id = pid
+
+        done = False
+        turn = 0
+        print("\n--- Stratego LLM Match Started ---")
+        print(f"Player 1 Agent: {agents[0].model_name}")
+        print(f"Player 2 Agent: {agents[1].model_name}")
+        if args.max_turns:
+            print(f"⏱️  Max turns limit: {args.max_turns} (testing mode)")
+        print()
+        while not done:
+            # Check max turns limit
+            if args.max_turns and turn >= args.max_turns:
+                print(f"\n⏱️  Reached max turns limit ({args.max_turns}). Stopping game early.")
+                break
+            
+            player_id, observation = env.get_observation()
+            current_agent = agents[player_id]
+            player_display = f"Player {player_id+1}"
+            model_name = current_agent.model_name
+            
+            # --- NEW LOGGING FOR TURN, PLAYER, AND MODEL ---
+            print(f"\n>>>> TURN {turn}: {player_display} ({model_name}) is moving...")
+
+            if (args.size == 10):
+                print_board(observation)
+            else:
+                print_board(observation, args.size)
+            # Pass recent move history to agent
+            current_agent.set_move_history(move_history[player_id][-10:])
+            history_str = tracker.to_prompt_string(player_id)
+            
+            # --- [CHANGE] INJECT AGGRESSION WARNING ---
+            # If the game drags on (e.g. > 20 turns), force them to wake up
+            if turn > 20:
+                observation += "\n\n[SYSTEM MESSAGE]: The game is stalling. You MUST ATTACK or ADVANCE immediately. Passive play is forbidden."
+            
+            if turn > 50:
+                 observation += "\n[CRITICAL]: STOP MOVING BACK AND FORTH. Pick a piece and move it FORWARD now."
+            # ------------------------------------------
+
+            observation = observation + history_str
+            # print(tracker.to_prompt_string(player_id))
+            lines = history_str.strip().splitlines()
+            if len(lines) <= 1:
+                print(history_str)
+            else:
+                header = lines[0:1]
+                body = lines[1:]
+                tail = body[-5:]  # Show only last 5 moves
+                print("\n".join(header + tail))
+            
+            # The agent (LLM) generates the action, retry a few times; fallback to available moves
+            action = ""
+            max_agent_attempts = 3
+            for attempt in range(max_agent_attempts):
+                action = current_agent(observation)
+                if action:
+                    break
+                print(f"[TURN {turn}] {model_name} failed to produce a move (attempt {attempt+1}/{max_agent_attempts}). Retrying...")
+
+            if not action:
+                legal = extract_legal_moves(observation)
+                forbidden = set(extract_forbidden(observation))
+                legal_filtered = [m for m in legal if m not in forbidden] or legal
+                if legal_filtered:
+                    action = random.choice(legal_filtered)
+                    print(f"[TURN {turn}] Fallback to random available move: {action}")
+                else:
+                    print(f"[TURN {turn}] No legal moves available for fallback; ending game loop.")
+                    break
+                
+            # --- NEW LOGGING FOR STRATEGY/MODEL DECISION ---
+            print(f"  > AGENT DECISION: {model_name} -> {action}")
+            print(f"  > Strategy/Model: Ollama Agent (T={current_agent.temperature}, Prompt='{args.prompt}')")
+
+            # Extract move details for logging
+            move_pattern = r'\[([A-J]\d+)\s+([A-J]\d+)\]'
+            match = re.search(move_pattern, action)
+            # src_pos = match.group(1) if match else ""
+            # dst_pos = match.group(2) if match else ""
+            
+            # # Get piece type from board (simplified extraction)
+            # piece_type = ""
+            # if src_pos and hasattr(env, 'game_state') and hasattr(env.game_state, 'board'):
+            #     try:
+            #         # Parse position like "D4" -> row=3, col=3
+            #         col = ord(src_pos[0]) - ord('A')
+            #         row = int(src_pos[1:]) - 1
+            #         piece = env.game_state.board[row][col]
+            #         if piece and hasattr(piece, 'rank_name'):
+            #             piece_type = piece.rank_name
+            #     except:
+            #         piece_type = "Unknown"
+            
+            # # Check if this is a repeated move (last 3 moves)
+            # was_repeated = False
+            # recent_moves = [m["move"] for m in move_history[player_id][-3:]]
+            # if action in recent_moves:
+            #     was_repeated = True
+            
+            # Record this move in history
+            move_history[player_id].append({
+                "turn": turn,
+                "move": action,
+                "text": f"Turn {turn}: You played {action}"
+            })
+
+            # Process move details for logging BEFORE making the environment step
+            move_details = process_move(
+                action=action,
+                board=env.env.board,
+                observation=observation,
+                player_id=player_id
+            )
+
+            # Execute the action exactly once in the environment
+            done, info = env.step(action=action)
+            
+            # Determine battle outcome by checking if target piece was there
+            battle_outcome = ""
+            if move_details.target_piece:
+                # There was a piece at destination, so battle occurred
+                # Check what's at destination now to determine outcome
+                dst_row = ord(move_details.dst_pos[0]) - ord('A')
+                dst_col = int(move_details.dst_pos[1:])
+                cell_after = env.env.board[dst_row][dst_col]
+                
+                if cell_after is None:
+                    # Both pieces removed = draw
+                    battle_outcome = "draw"
+                elif isinstance(cell_after, dict):
+                    if cell_after.get('player') == player_id:
+                        battle_outcome = "won"
+                    else:
+                        battle_outcome = "lost"
+            
+            # Extract outcome from environment observation
+            outcome = "move"
+            # captured = ""
+            obs_text = ""
+            # if isinstance(info, (list, tuple)) and len(info) > 1:
+            #     obs_text = str(info[1])
+            # else:
+            #     obs_text = str(info)
+            if isinstance(info, (list, tuple)):
+                if 0 <= player_id < len(info):
+                    obs_text = str(info[player_id])
+                else:
+                    obs_text = " ".join(str(x) for x in info)
+            else:
+                obs_text = str(info)
+
+            low = obs_text.lower()
+            if "invalid" in low or "illegal" in low:
+                outcome = "invalid"
+            elif "captured" in low or "won the battle" in low:
+                outcome = "won_battle"
+            elif "lost the battle" in low or "defeated" in low:
+                outcome = "lost_battle"
+            elif "draw" in low or "tie" in low:
+                outcome = "draw"
+                    
+            event = info.get("event") if isinstance(info, dict) else None
+            extra = info.get("detail") if isinstance(info, dict) else None
+            
+            if outcome != "invalid":
+                # Record this move in history
+                move_history[player_id].append({
+                    "turn": turn,
+                    "move": action,
+                    "text": f"Turn {turn}: You played {action}"
+                })
+
+                tracker.record(
+                    player=player_id,
+                    move=action,
+                    event=event,
+                    extra=extra
+                )
+            else:
+                move_history[player_id].append({
+                    "turn": turn,
+                    "move": action,
+                    "text": f"Turn {turn}: INVALID move {action}"
+                })
+                tracker.record(
+                    player=player_id,
+                    move=action,
+                    event="invalid_move",
+                    extra=extra
+                )
+                print(f"[HISTORY] Skipping invalid move from history: {action}")
+
+            logger.log_move(turn=turn,
+                                player=player_id,
+                                model_name=getattr(current_agent, "model_name", "unknown"),
+                                move=action,
+                                src=move_details.src_pos,
+                                dst=move_details.dst_pos,
+                                piece_type=move_details.piece_type,
+                                board_state=move_details.board_state,
+                                available_moves=move_details.available_moves,
+                                move_direction=move_details.move_direction,
+                                target_piece=move_details.target_piece,
+                                battle_outcome=battle_outcome,
+                            )
+            turn += 1
+
+
+        # --- Game Over & Winner Announcement ---
+        rewards, game_info = env.close()
+        print("\n" + "="*50)
+        print("--- GAME OVER ---")
+        game_duration = time.time() - game_start_time
+        # Print summary
+        print(f"\nGame finished. Duration: {int(game_duration // 60)}m {int(game_duration % 60)}s")
+        print(f"Result: {rewards} | {game_info}")
+        
+        # Logic to declare the specific winner based on rewards
+        # Rewards are usually {0: 1, 1: -1} (P0 Wins) or {0: -1, 1: 1} (P1 Wins)
+        p0_score = rewards.get(0, 0)
+        p1_score = rewards.get(1, 0)
+        winner = None
+        game_result = ""
+
+        if p0_score > p1_score:
+            winner = 0
+            game_result = "player0"
+            print(f"\n🏆 * * * PLAYER 0 WINS! * * * 🏆")
+            print(f"Agent: {agents[0].model_name}")
+        elif p1_score > p0_score:
+            winner = 1
+            game_result = "player1"
+            print(f"\n🏆 * * * PLAYER 1 WINS! * * * 🏆")
+            print(f"Agent: {agents[1].model_name}")
+        else:
+            game_result = "draw"
+            print(f"\n🤝 * * * IT'S A DRAW! * * * 🤝")
+
+        print("\nDetails:")
+        print(f"Final Rewards: {rewards}")
+        print(f"Game Info: {game_info}")
+    
+    try:
+        invalid_players = [
+            pid for pid, info_dict in (game_info or {}).items()
+            if isinstance(info_dict, dict) and info_dict.get("invalid_move")
+        ]
+        if invalid_players:
+            import csv
+            csv_path = logger.path
+            rows = []
+            fieldnames = None
+
+            with open(csv_path, "r", encoding="utf-8", newline="") as f:
+                reader = csv.DictReader(f)
+                fieldnames = reader.fieldnames
+                for r in reader:
+                    rows.append(r)
+
+            if rows and fieldnames and "outcome" in fieldnames:
+                rows[-1]["outcome"] = "invalid"
+                with open(csv_path, "w", encoding="utf-8", newline="") as f:
+                    writer = csv.DictWriter(f, fieldnames=fieldnames)
+                    writer.writeheader()
+                    writer.writerows(rows)
+
+                print("\n[LOG PATCH] Last move outcome patched to 'invalid' "
+                      f"(player {invalid_players[0]} made an invalid move).")
+                
+    except Exception as e:
+        print(f"[LOG PATCH] Failed to patch CSV outcome: {e}")
+            
+        # Finalize the game log with winner info in every row
+        logger.finalize_game(winner=winner, game_result=game_result)
+    
+    # LLM analyzes the game CSV and updates prompt
+    analyze_and_update_prompt(
+        csv_path=logger.path,
+        prompts_dir="stratego/prompts",
+        logs_dir=args.log_dir,
+        model_name="mistral:7b",  # Analysis model
+        models_used=[agents[0].model_name, agents[1].model_name],
+        game_duration_seconds=game_duration,
+        winner=winner,
+        total_turns=turn - 1
+    )
+    
+    # Auto-push game data to Hugging Face Hub
+    print("\nSyncing game data to Hugging Face...")
+    auto_push_after_game(
+        logs_dir=os.path.join(args.log_dir, "games"),
+        repo_id="STRATEGO-LLM-TRAINING/stratego",
+    )