test.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:
            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


if __name__ == "__main__":
    # Create the environment
    vec_env = DummyVecEnv([lambda: RPGvsRPGEnv()])

    # Initially, there is no opponent model, so the agent trains against a random one.
    # Load a pre-existing model or start fresh
    if os.path.exists(f"{MODEL_PATH}.zip"):
        print("Loading existing model...")
        model = PPO.load(MODEL_PATH, env=vec_env)
    else:
        print("No existing model found. Starting new training.")
        model = PPO("MlpPolicy", vec_env, verbose=1)

    # Create the callback
    callback = OpponentUpdateCallback(MODEL_PATH, OPPONENT_PATH, OPPONENT_UPDATE_FREQ)

    print("\nStarting self-play training...")
    # The callback will handle updating the opponent during training
    model.learn(total_timesteps=TRAINING_TIMESTEPS, callback=callback)

    print("\nTraining finished.")
    model.save(MODEL_PATH)

    # --- Evaluation ---
    print("\nEvaluating final agent against its last opponent...")
    eval_env = RPGvsRPGEnv()
    opponent = PPO.load(OPPONENT_PATH, env=eval_env)
    eval_env.set_opponent_model(opponent)

    wins = 0
    losses = 0
    draws = 0
    episodes = 100
    
    obs, _ = eval_env.reset()
    for i in range(episodes):
        action, _ = model.predict(obs, deterministic=True)
        print(action)
        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%})")
    print(f"Losses: {losses} ({losses/episodes:.2%})")
    print(f"Draws: {draws} ({draws/episodes:.2%})")