test2.py

import gymnasium as gym
import numpy as np
from gymnasium import spaces
import random
import os
from stable_baselines3 import PPO
from stable_baselines3.common.env_checker import check_env
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.callbacks import BaseCallback

# --- Configuration ---
TRAINING_TIMESTEPS = 200_000
OPPONENT_UPDATE_FREQ = 25_000 # How often to update the opponent model
MODEL_PATH = "rpg_self_play_ppo"
OPPONENT_PATH = "rpg_opponent_ppo.zip"


class Player:
    """A helper class to manage player state."""
    def __init__(self, start_money):
        self.start_money = start_money
        self.reset()

    def reset(self):
        self.hp = 100
        self.max_hp = 100
        self.base_dmg = 10
        self.bonus_dmg = 0
        self.money = self.start_money
        self.inventory = []

class RPGvsRPGEnv(gym.Env):
    """
    A self-play RPG environment where an RL agent fights another RL agent.
    The opponent's purchases are unknown to the player (imperfect information).
    """
    def __init__(self, opponent_model=None):
        super(RPGvsRPGEnv, self).__init__()
        self.start_money = 700
        
        # The agent being trained
        self.player = Player(self.start_money)
        # The opponent agent
        self.opponent = Player(self.start_money)

        # The model for the opponent agent
        self.opponent_model = opponent_model

        # Action space: [Armour, Item1, Item2]
        self.action_space = spaces.MultiDiscrete([4, 6, 6])

        # Observation space: [MyHP, MyMoney, OpponentHP]
        # Imperfect information: agent does not see opponent's money or purchases.
        self.observation_space = spaces.Box(
            low=np.array([0, 0, 0]),
            high=np.array([200, self.start_money, 200]), # MaxHP, MaxMoney, MaxOpponentHP
            dtype=np.float32
        )

        # Game data
        self.armours = [
            {"name": "None", "hp": 0, "cost": 0}, {"name": "Leather", "hp": 20, "cost": 200},
            {"name": "Chainmail", "hp": 40, "cost": 400}, {"name": "Plate", "hp": 60, "cost": 600},
        ]
        self.items = [
            None, {"name": "Small Heal", "heal": 30, "cost": 100}, {"name": "Large Heal", "heal": 60, "cost": 200},
            {"name": "Damage Boost", "bonus": 10, "cost": 150}, {"name": "SP Boost", "sp": 50, "cost": 150},
            {"name": "Max SP", "spmax": 50, "cost": 200},
        ]

    def set_opponent_model(self, opponent_model):
        """Method to update the opponent model during training."""
        self.opponent_model = opponent_model

    def _apply_actions_to_player(self, player_obj, action):
        """Helper to apply purchase actions to a player object."""
        armour_choice, item1_choice, item2_choice = action
        
        # Buy Armour
        if armour_choice > 0:
            armour = self.armours[armour_choice]
            if player_obj.money >= armour["cost"]:
                player_obj.max_hp = 100 + armour["hp"]
                player_obj.hp = player_obj.max_hp
                player_obj.money -= armour["cost"]

        # Buy Items
        for item_choice in [item1_choice, item2_choice]:
            if item_choice > 0:
                item = self.items[item_choice]
                if item and player_obj.money >= item["cost"]:
                    player_obj.inventory.append(item)
                    player_obj.money -= item["cost"]
        
        # Use Items
        for item in player_obj.inventory:
            if "heal" in item:
                player_obj.hp += item["heal"]
            if "bonus" in item:
                player_obj.bonus_dmg += item["bonus"]
        player_obj.inventory.clear()

    def reset(self, *, seed=None, options=None):
        super().reset(seed=seed)
        self.player.reset()
        self.opponent.reset()
        return self._get_obs(self.player, self.opponent), {}

    def _get_obs(self, primary_player, other_player):
        """Gets observation from the perspective of the primary_player."""
        return np.array([
            max(0, primary_player.hp),
            primary_player.money,
            max(0, other_player.hp)
        ], dtype=np.float32)

    def step(self, action):
        # 1. Get opponent's action
        if self.opponent_model:
            opponent_obs = self._get_obs(self.opponent, self.player)
            opponent_action, _ = self.opponent_model.predict(opponent_obs, deterministic=True)
        else:
            # If no model, opponent does nothing (or takes a random action)
            opponent_action = self.action_space.sample()

        # 2. Apply actions to both players
        self._apply_actions_to_player(self.player, action)
        self._apply_actions_to_player(self.opponent, opponent_action)

        # 3. Combat Simulation
        while self.player.hp > 0 and self.opponent.hp > 0:
            # Player attacks opponent
            self.opponent.hp -= (self.player.base_dmg + self.player.bonus_dmg)
            if self.opponent.hp <= 0:
                break
            # Opponent attacks player
            self.player.hp -= (self.opponent.base_dmg + self.opponent.bonus_dmg)

        # 4. Determine reward and done status
        if self.player.hp > 0 and self.opponent.hp <= 0:
            reward = 1.0  # Player wins
        elif self.player.hp <= 0 and self.opponent.hp > 0:
            reward = -1.0 # Player loses
        else:
            reward = 0.0 # Draw

        done = True
        return self._get_obs(self.player, self.opponent), reward, done, False, {}


class OpponentUpdateCallback(BaseCallback):
    """
    A custom callback to update the opponent model during training.
    """
    def __init__(self, model_path, opponent_path, update_freq, verbose=0):
        super(OpponentUpdateCallback, self).__init__(verbose)
        self.model_path = model_path
        self.opponent_path = opponent_path
        self.update_freq = update_freq

    def _on_step(self) -> bool:
        if self.n_calls % self.update_freq == 0 and self.n_calls > 0:
            print(f"\nCallback called at step {self.n_calls}. Updating opponent...")
            # Save the current model
            self.model.save(self.model_path)
            # Copy it to become the new opponent model
            # In a real scenario, you'd copy the file. Here we just save/load.
            self.model.save(self.opponent_path)
            
            # Load the new opponent model
            new_opponent_model = PPO.load(self.opponent_path, env=self.training_env)
            
            # Update the opponent model in the vectorized environment
            self.training_env.env_method("set_opponent_model", new_opponent_model)
            print("Opponent updated successfully.")
        return True

def get_human_action(money, armours, items):
    """Function to handle human player input for purchasing."""
    print("\n--- YOUR TURN TO BUY ---")
    print(f"You have ${money}.")
    
    # Display armour options
    print("\n--- Armours ---")
    for i, armour in enumerate(armours):
        print(f"[{i}] {armour['name']:<10} | HP: +{armour['hp']:<3} | Cost: ${armour['cost']}")
    
    # Get armour choice
    while True:
        try:
            armour_choice = int(input("Choose your armour [0-3]: "))
            if 0 <= armour_choice < len(armours) and money >= armours[armour_choice]['cost']:
                money -= armours[armour_choice]['cost']
                break
            elif 0 <= armour_choice < len(armours):
                print("Not enough money for that armour!")
            else:
                print("Invalid choice.")
        except ValueError:
            print("Invalid input. Please enter a number.")

    # Display item options
    print(f"\nRemaining money: ${money}")
    print("\n--- Items ---")
    for i, item in enumerate(items):
        if item:
            print(f"[{i}] {item['name']:<12} | Cost: ${item['cost']}")
        else:
            print(f"[0] None")
            
    # Get item 1 choice
    item1_choice = 0
    while True:
        try:
            item1_choice = int(input("Choose your first item [0-5]: "))
            if 0 <= item1_choice < len(items) and (item1_choice == 0 or money >= items[item1_choice]['cost']):
                if item1_choice > 0:
                    money -= items[item1_choice]['cost']
                break
            elif 0 <= item1_choice < len(items):
                print("Not enough money for that item!")
            else:
                print("Invalid choice.")
        except ValueError:
            print("Invalid input. Please enter a number.")
    
    # Get item 2 choice
    item2_choice = 0
    print(f"\nRemaining money: ${money}")
    while True:
        try:
            item2_choice = int(input("Choose your second item [0-5]: "))
            if 0 <= item2_choice < len(items) and (item2_choice == 0 or money >= items[item2_choice]['cost']):
                break
            elif 0 <= item2_choice < len(items):
                print("Not enough money for that item!")
            else:
                print("Invalid choice.")
        except ValueError:
            print("Invalid input. Please enter a number.")
            
    return [armour_choice, item1_choice, item2_choice]

def play_vs_ai_game(model):
    """Main loop for a human to play against the trained AI model."""
    env = RPGvsRPGEnv()
    
    while True:
        human_player = Player(env.start_money)
        ai_player = Player(env.start_money)

        # Get human action
        human_action = get_human_action(human_player.money, env.armours, env.items)
        env._apply_actions_to_player(human_player, human_action)
        
        # Get AI action
        ai_obs = env._get_obs(ai_player, human_player)
        ai_action, _ = model.predict(ai_obs, deterministic=True)
        env._apply_actions_to_player(ai_player, ai_action)

        print("\n--- PRE-BATTLE STATS ---")
        print(f"Your HP: {human_player.hp}/{human_player.max_hp} | Your Bonus Dmg: {human_player.bonus_dmg}")
        ai_armour = env.armours[ai_action[0]]['name']
        ai_item1 = env.items[ai_action[1]]['name'] if ai_action[1] > 0 else "None"
        ai_item2 = env.items[ai_action[2]]['name'] if ai_action[2] > 0 else "None"
        print(f"AI bought: {ai_armour}, {ai_item1}, {ai_item2}")
        print(f"AI HP: {ai_player.hp}/{ai_player.max_hp} | AI Bonus Dmg: {ai_player.bonus_dmg}")
        print("\n--- COMBAT START ---")

        # Combat
        turn = 1
        while human_player.hp > 0 and ai_player.hp > 0:
            print(f"\n- Turn {turn} -")
            # Human attacks
            ai_player.hp -= (human_player.base_dmg + human_player.bonus_dmg)
            print(f"You attack the AI! AI HP is now {max(0, ai_player.hp)}.")
            if ai_player.hp <= 0: break

            # AI attacks
            human_player.hp -= (ai_player.base_dmg + ai_player.bonus_dmg)
            print(f"The AI attacks you! Your HP is now {max(0, human_player.hp)}.")
            turn += 1

        print("\n--- COMBAT END ---")
        if human_player.hp > 0 and ai_player.hp <= 0:
            print("🎉 You are victorious! �")
        elif human_player.hp <= 0 and ai_player.hp > 0:
            print("☠️ You have been defeated. ☠️")
        else:
            print("🤝 The battle is a draw! 🤝")
        
        play_again = input("\nPlay again? (y/n): ").lower()
        if play_again != 'y':
            break

if __name__ == "__main__":
    
    while True:
        print("\n--- RPG Self-Play AI ---")
        print("[1] Train model")
        print("[2] Evaluate model vs. its past self")
        print("[3] Play against the trained AI")
        print("[4] Exit")
        choice = input("Enter your choice: ")

        if choice == '1':
            vec_env = DummyVecEnv([lambda: RPGvsRPGEnv()])
            if os.path.exists(f"{MODEL_PATH}.zip"):
                print("Loading existing model to continue training...")
                model = PPO.load(MODEL_PATH, env=vec_env)
            else:
                print("No existing model found. Starting new training.")
                model = PPO("MlpPolicy", vec_env, verbose=1)
            callback = OpponentUpdateCallback(MODEL_PATH, OPPONENT_PATH, OPPONENT_UPDATE_FREQ)
            print("\nStarting self-play training...")
            model.learn(total_timesteps=TRAINING_TIMESTEPS, callback=callback)
            print("\nTraining finished.")
            model.save(MODEL_PATH)

        elif choice == '2':
            if not os.path.exists(OPPONENT_PATH):
                print("No opponent model found! Please train the model first (Option 1).")
                continue
            print("\nEvaluating final agent against its last opponent...")
            model = PPO.load(MODEL_PATH)
            eval_env = RPGvsRPGEnv()
            opponent = PPO.load(OPPONENT_PATH, env=eval_env)
            eval_env.set_opponent_model(opponent)
            wins, losses, draws, episodes = 0, 0, 0, 100
            obs, _ = eval_env.reset()
            for i in range(episodes):
                action, _ = model.predict(obs, deterministic=True)
                obs, reward, done, _, info = eval_env.step(action)
                if reward > 0: wins += 1
                elif reward < 0: losses += 1
                else: draws += 1
                obs, _ = eval_env.reset()
            print(f"Evaluation over {episodes} episodes:")
            print(f"Wins: {wins} ({wins/episodes:.2%}) | Losses: {losses} ({losses/episodes:.2%}) | Draws: {draws} ({draws/episodes:.2%})")

        elif choice == '3':
            if not os.path.exists(f"{MODEL_PATH}.zip"):
                print("No trained model found! Please train the model first (Option 1).")
                continue
            model = PPO.load(MODEL_PATH)
            play_vs_ai_game(model)

        elif choice == '4':
            print("Exiting.")
            break
        else:
            print("Invalid choice. Please try again.")