src/neural_ppo_old.py - master - GraphTreeHeap/NeuroFighter

src/neural_ppo_old.py
315 строк13 KB
FedMam
moving to the new project version
11 окт 2025, 04:16
11 окт 2025, 04:1600409b0
import torch
import torch.nn as nn
import numpy as np
import copy
import random
import tqdm
from collections import deque
from mechanics import *
from robotics import *
from typing import *


class PPOAgentNoEmbeddings(nn.Module):
    def __init__(self,
                 state_dim,
                 action_dim,
                 n_characters,
                 character_emb_dim=32,
                 hidden_dim=256):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, hidden_dim),
            nn.LeakyReLU(0.2),
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
        )
        self.actor = nn.Linear(hidden_dim, action_dim)
        self.critic = nn.Linear(hidden_dim, 1)

    def forward(self, my_character, opp_character, state):
        state = state.squeeze()
        features = self.backbone(state)
        logits = self.actor(features)
        value = self.critic(features).squeeze()
        return logits, value


class PPOAgent(nn.Module):
    def __init__(self,
                 state_dim,
                 action_dim,
                 n_characters,
                 character_emb_dim=32,
                 hidden_dim=256):
        super().__init__()
        self.my_character_emb = nn.Embedding(n_characters, character_emb_dim)
        self.opp_character_emb = nn.Embedding(n_characters, character_emb_dim)
        self.backbone = nn.Sequential(
            nn.Linear(character_emb_dim * 2 + state_dim, hidden_dim),
            nn.ReLU(),
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, hidden_dim),
            nn.LeakyReLU(0.2),
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
        )
        self.actor = nn.Linear(hidden_dim, action_dim)
        self.critic = nn.Linear(hidden_dim, 1)

    def forward(self, my_character, opp_character, state):
        my_character = self.my_character_emb(my_character).squeeze()
        opp_character = self.opp_character_emb(opp_character).squeeze()
        state = torch.cat((my_character, opp_character, state.squeeze()), dim=-1)
        features = self.backbone(state)
        logits = self.actor(features)
        value = self.critic(features).squeeze()
        return logits, value
    

def collect_trajectories_ppo(current_agent: PPOAgent,
                        opponents: Iterable[Any],
                        temperature: float,
                        num_episodes: int=10,
                        rand: random.Random | None=None):
    if rand is None:
        rand = random.Random()

    my_characters, opp_characters, states, actions, rewards, dones, next_states, old_log_probs = [], [], [], [], [], [], [], []
    
    for _ in range(num_episodes):
        env = Environment(player_desc=CHARACTERS[rand.randint(0, len(CHARACTERS)-1)],
                          opponent_desc=CHARACTERS[rand.randint(0, len(CHARACTERS)-1)])
        agent_is_player = rand.random() < 0.5

        opponent = rand.choice(opponents)
        
        
        while not env.done():
            state_agent = env.get_state(player=agent_is_player)
            state_opponent = env.get_state(player=not agent_is_player)
            agent_character = env.player.character_id if agent_is_player else env.opponent.character_id
            opponent_character = env.opponent.character_id if agent_is_player else env.player.character_id
            
            # Current agent acts
            logits, _ = current_agent(torch.LongTensor((agent_character,)),
                                      torch.LongTensor((opponent_character,)),
                                      torch.FloatTensor(state_agent))
            probs = torch.softmax(logits / temperature, dim=-1)
            dist = torch.distributions.Categorical(probs)
            action = dist.sample()
            old_log_prob = dist.log_prob(action)
            env.take_action(agent_is_player, action.item())
            
            # Opponent acts
            if isinstance(opponent, Bot):
                opponent.decide_outer(env, not agent_is_player)
            else:
                with torch.no_grad():
                    opp_logits, _ = opponent(torch.LongTensor((opponent_character,)),
                                            torch.LongTensor((agent_character,)),
                                            torch.FloatTensor(state_opponent))
                    opp_probs = torch.softmax(opp_logits / temperature, dim=-1)
                    opp_dist = torch.distributions.Categorical(opp_probs)
                    opp_action = opp_dist.sample()
                    env.take_action(not agent_is_player, opp_action.item())
            
            # Step environment
            player_reward, opponent_reward = env.step()
            reward = (player_reward if agent_is_player else opponent_reward)
            next_state_agent = env.get_state(player=agent_is_player)
            done = env.done()
            
            my_characters.append(agent_character)
            opp_characters.append(opponent_character)
            states.append(state_agent)
            actions.append(action.item())
            rewards.append(reward)
            dones.append(done)
            next_states.append(next_state_agent)
            old_log_probs.append(old_log_prob.item())
    
    return {
        "my_characters": torch.LongTensor(my_characters),
        "opponent_characters": torch.LongTensor(opp_characters),
        "states": torch.FloatTensor(states),
        "actions": torch.LongTensor(actions),
        "rewards": torch.FloatTensor(rewards),
        "dones": torch.FloatTensor(dones),
        "next_states": torch.FloatTensor(next_states),
        "old_log_probs": torch.FloatTensor(old_log_probs)
    }


class PPOAgentWrapper:
    def __init__(self,
                 character_emb_dim: int=32,
                 hidden_dim: int=256,
                 temperature: float=1.0,
                 entropy_coef_start: float=0.02,
                 entropy_coef_min: float=0.0,
                 entropy_coef_decay_episodes: int | None=500,
                 agent_class: type[Any]=PPOAgent):
        self.character_emb_dim = character_emb_dim
        self.hidden_dim = hidden_dim
        self.agent_class = agent_class
        self.current_agent = agent_class(STATE_DIM, ACTION_DIM_S, len(CHARACTERS), character_emb_dim, hidden_dim)
        self.optimizer = torch.optim.Adam(self.current_agent.parameters(), lr=3e-4)

        self.temperature = temperature
        self.entropy_coef = entropy_coef_start
        self.entropy_coef_min = entropy_coef_min
        self.entropy_coef_decay = 0 if entropy_coef_decay_episodes is None else (entropy_coef_start - entropy_coef_min) / entropy_coef_decay_episodes

    def load_model(self, filename):
        self.current_agent.load_state_dict(torch.load(filename))

    def update(self,
            data: dict[str, torch.Tensor],
            gamma=0.99,
            gae_lambda=0.95,
            clip_eps=0.2,
            epochs=4,
            batch_size=64):
        my_characters, opp_characters, states, actions, rewards, dones, next_states, old_log_probs = data.values()
        
        # Compute advantages with GAE
        with torch.no_grad():
            _, values = self.current_agent(my_characters, opp_characters, states)
            _, next_values = self.current_agent(my_characters, opp_characters, next_states)
            deltas = rewards + gamma * next_values * (1 - dones) - values
            advantages = torch.zeros_like(rewards)
            advantage = 0
            for t in reversed(range(len(rewards))):
                advantage = deltas[t] + gamma * gae_lambda * (1 - dones[t]) * advantage
                advantages[t] = advantage
            returns = advantages + values
        
        # Normalize advantages
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        
        # PPO epochs
        for _ in range(epochs):
            indices = torch.randperm(len(states))
            for i in range(0, len(states), batch_size):
                idx = indices[i:i+batch_size]
                batch = {k: v[idx] for k, v in data.items()}
                
                # Evaluate current policy
                logits, values = self.current_agent(batch["my_characters"], batch["opponent_characters"], batch["states"])
                probs = torch.softmax(logits / self.temperature, dim=-1)
                dist = torch.distributions.Categorical(probs)
                log_probs = dist.log_prob(batch["actions"])
                entropy = dist.entropy().mean()
                
                # PPO loss
                ratios = torch.exp(log_probs - batch["old_log_probs"])
                surr1 = ratios * advantages[idx]
                surr2 = torch.clamp(ratios, 1-clip_eps, 1+clip_eps) * advantages[idx]
                actor_loss = -torch.min(surr1, surr2).mean()
                
                # Critic loss
                critic_loss = (values - returns[idx]).pow(2).mean()
                
                # Total loss
                loss = actor_loss + 0.5 * critic_loss - self.entropy_coef * entropy
                
                # Update
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

    def train(self,
            opponents: Iterable[Any],
            iterations: int=1000,
            save_state_dict_freq: int | None=50,
            save_state_dict_name: str='trained_agent',
            verbose: bool=False,
            experiment_id: int | None=None,
            rand: random.Random | None=None):
        rand = rand or random.Random()
        experiment_id = experiment_id or random.randint(0, 99999)
        if verbose:
            print('Experiment ID:', experiment_id)

        self.current_agent.train()

        for iter in tqdm.tqdm(range(iterations), 'Training PPO agent') if verbose else range(iterations):
            # Collect data against random opponent from pool
            data = collect_trajectories_ppo(self.current_agent, opponents, temperature=self.temperature, num_episodes=10, rand=rand)
            
            # Update current agent
            self.update(data)

            # decay temperature & entropy coef (linear decay)
            self.temperature = max(self.temperature_min, self.temperature - self.temperature_decay)
            self.entropy_coef = max(self.entropy_coef_min, self.entropy_coef - self.entropy_coef_decay)
            
            if save_state_dict_freq is not None and ((iter + 1) % save_state_dict_freq == 0 or iter == iterations - 1):
                torch.save(self.current_agent.state_dict(), f'{save_state_dict_name}_{experiment_id}.pth')
    
    def train_self_play(self,
                        epochs: int=100,
                        iterations_per_epoch: int=10,
                        opponent_pool_size: int=5,
                        save_state_dict_name='trained_agent',
                        experiment_id: int | None=None,
                        verbose: bool=False,
                        rand: random.Random | None = None):
        rand = rand or random.Random()
        experiment_id = experiment_id or random.randint(0, 99999)
        if verbose:
            print('Experiment ID:', experiment_id)

        self.current_agent.train()

        opponent_pool = deque(maxlen=opponent_pool_size)

        for epoch in tqdm.tqdm(range(epochs), 'Training self-play PPO agent') if verbose else range(epochs):
            old_version = self.agent_class(STATE_DIM, ACTION_DIM_S, len(CHARACTERS), self.character_emb_dim, self.hidden_dim)
            old_version.load_state_dict(copy.deepcopy(self.current_agent.state_dict()))
            old_version.eval()
            opponent_pool.append(old_version)

            self.train(opponent_pool,
                       iterations=iterations_per_epoch,
                       save_state_dict_freq=None,
                       save_state_dict_name=save_state_dict_name,
                       experiment_id=experiment_id,
                       verbose=False,
                       rand=rand)
            
            torch.save(self.current_agent.state_dict(), f'{save_state_dict_name}_{experiment_id}.pth')


def play_ppo(ppo_agent: PPOAgentWrapper,
         env: Environment):
    import graphics

    ppo_agent.current_agent.eval()

    graphics.init_graphics()
    tv = graphics.GameTV(env, player_controlled=True)

    while not env.done():
        state_agent = torch.FloatTensor(env.get_state(player=False))
        logits, _ = ppo_agent.current_agent(torch.LongTensor((env.opponent.character_id,)),
                                            torch.LongTensor((env.player.character_id,)),
                                            torch.FloatTensor(state_agent))
        probs = torch.softmax(logits / ppo_agent.temperature, dim=-1)
        dist = torch.distributions.Categorical(probs)
        action = dist.sample()
        env.take_action(False, action.item())

        # DEBUG
        # print('~', probs)

        env.step()
        tv.step()
    
    graphics.deinit_graphics()