import numpy as np
import torch
import gymnasium as gym
from ppo_torch import Agent
import time, telegram, asyncio
from wrappers import make_atari, wrap_deepmind
from gymnasium.vector import SyncVectorEnv

channel_id = "-1002031363666"
bot = telegram.Bot(token="<TRUFFLEHOG_REDACTED_TELEGRAMBOTTOKEN>")

async def send_message_to_channel(message_text):
    await bot.send_message(channel_id, message_text)



import numpy as np
import torch as T
from ppo_memory import PPOMemory
from network import CriticConvNetwork_v2, ActorConvNetwork_v2

class Agent:
    def __init__(self, n_actions, input_dims, gamma=0.99, alpha=0.0003, gae_lambda=0.95,
                 policy_clip=0.2, batch_size=64, n_epochs=10, import_models = False):
        self.gamma = gamma
        self.policy_clip = policy_clip
        self.n_epochs = n_epochs
        self.gae_lambda = gae_lambda

        self.actor = ActorConvNetwork_v2(n_actions, input_dims, alpha)
        self.critic = CriticConvNetwork_v2(input_dims, alpha)

        if import_models:
            self.actor.load_checkpoint()
            self.critic.load_checkpoint()

        self.memory = PPOMemory(batch_size)

    def remember(self, state, action, probs, vals, reward, done):
        self.memory.store_memory(state, action, probs, vals, reward, done)

    def save_models(self):
        print('... saving models ...')
        self.actor.save_checkpoint()
        self.critic.save_checkpoint()

    def load_models(self):
        print('... loading models ...')
        self.actor.load_checkpoint()
        self.critic.load_checkpoint()

    @T.no_grad()
    def choose_action(self, observations):
        # observations shape: [num_envs, channels, height, width]
        states = T.tensor(observations, dtype=T.float).to(self.actor.device)

        dist = self.actor(states)
        values = self.critic(states)
        actions = dist.sample()

        probs = dist.log_prob(actions)
        # No need to squeeze or convert to items; keep tensors as-is
        # actions shape: [num_envs]
        # probs shape: [num_envs]
        # values shape: [num_envs]

        return actions.cpu().numpy(), probs.cpu().detach().numpy(), values.squeeze().cpu().detach().numpy()

    def learn(self):
        for _ in range(self.n_epochs):
            state_arr, action_arr, old_prob_arr, vals_arr, \
                reward_arr, dones_arr, batches = \
                self.memory.generate_batches()

            vals = vals_arr.squeeze()

            advantages = np.zeros(len(reward_arr), dtype=np.float32)
            advantage = 0.0

            # Compute GAE using a single reversed loop
            for t in reversed(range(len(reward_arr)-1)):
                mask = 1.0 - dones_arr[t]
                delta = reward_arr[t] + self.gamma * vals[t + 1] * mask - vals[t]
                advantage = delta + self.gamma * self.gae_lambda * mask * advantage
                advantages[t] = advantage

            # Normalize advantages for better training stability
            advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
            advantages = T.tensor(advantages, dtype=T.float).to(self.actor.device)
            vals = T.tensor(vals_arr, dtype=T.float).squeeze().to(self.actor.device)

            for batch in batches:
                self.actor.optimizer.zero_grad()
                self.critic.optimizer.zero_grad()

                states = T.tensor(state_arr[batch], dtype=T.float).to(self.actor.device)
                old_probs = T.tensor(old_prob_arr[batch]).to(self.actor.device)
                actions = T.tensor(action_arr[batch]).to(self.actor.device)
                batch_advantages = advantages[batch]
                batch_vals = vals[batch]

                dist = self.actor(states)
                entropy = dist.entropy().mean()

                critic_value = self.critic(states).squeeze()

                new_probs = dist.log_prob(actions)
                prob_ratio = (new_probs - old_probs).exp()
                weighted_probs = batch_advantages * prob_ratio
                weighted_clipped_probs = T.clamp(prob_ratio, 1 - self.policy_clip,
                                                 1 + self.policy_clip) * batch_advantages
                actor_loss = -T.min(weighted_probs, weighted_clipped_probs).mean()

                returns = batch_advantages + batch_vals
                critic_loss = (returns - critic_value) ** 2
                critic_loss = critic_loss.mean()

                total_loss = actor_loss + (0.5 * critic_loss) - (0.01 * entropy)

                total_loss.backward()
                T.nn.utils.clip_grad_norm_(self.actor.parameters(), 0.5)
                self.actor.optimizer.step()
                self.critic.optimizer.step()
                self.actor.scheduler.step()
                self.critic.scheduler.step()

        self.memory.clear_memory()



import numpy as np
import torch, os
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from lion_pytorch import Lion
from torch.distributions.categorical import Categorical
from torch.optim.lr_scheduler import StepLR


class DeepQNetwork(nn.Module):
    def __init__(self, lr, input_dims, fc1_dims, fc2_dims, action_size):
        super(DeepQNetwork, self).__init__()
        self.input_dims = input_dims                            # Osservazioni dell'ambiente
        self.fc1_dims = fc1_dims                                # Neuroni del primo layer
        self.fc2_dims = fc2_dims                                # Neuroni del secondo layer
        self.action_size = action_size                          # Numero di azioni possibili
        self.fc1 = nn.Linear(*self.input_dims, self.fc1_dims)   # Primo Hidden Layer
        self.fc2 = nn.Linear(self.fc1_dims, self.fc2_dims)      # Secondo Hidden Layer
        self.adv_layer = nn.Linear(self.fc2_dims, action_size)     # Output Layer
        self.value_layer = nn.Linear(self.fc2_dims, 1)

        self.optimizer = optim.Adam(self.parameters(), lr=lr)   # Adam Optimizer
        self.loss = nn.MSELoss()                                # Mean Square Error Loss
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        self.to(self.device)

    def forward(self, state):
        x = F.relu(self.fc1(state))
        x = F.relu(self.fc2(x))
        adv = self.adv_layer(x)
        value = self.value_layer(x)

        return value, adv

class ConvQNet(nn.Module):
    def __init__(self, lr=1e-6, action_size=4):
        super(ConvQNet, self).__init__()
        self.conv1 = nn.Conv2d(4, 32, kernel_size=8, stride=4)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)

        self.bn = nn.BatchNorm2d(32)
        self.bn2 = nn.BatchNorm2d(64)
        self.bn3 = nn.BatchNorm2d(64)

        self.fc1 = nn.Linear(7 * 7 * 64, 512)
        self.adv = nn.Linear(512, action_size)
        self.value = nn.Linear(512, 1)

        self.optimizer = optim.Adam(self.parameters(), lr=lr)  # Adam Optimizer
        self.loss = nn.MSELoss()  # Mean Square Error Loss
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        self.to(self.device)

    def forward(self, state):
        x = F.relu(self.conv1(state))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        adv = F.relu(self.adv(x))
        value = F.relu(self.value(x))

        return value, adv


class ActorNetwork(nn.Module):
    def __init__(self, n_actions, input_dims, alpha,
                 fc1_dims=256, fc2_dims=256, chkpt_dir='tmp/ppo'):
        super(ActorNetwork, self).__init__()

        self.checkpoint_file = os.path.join(chkpt_dir, 'actor_torch_ppo')
        self.actor = nn.Sequential(
            nn.Linear(*input_dims, fc1_dims),
            nn.ReLU(),
            nn.Linear(fc1_dims, fc2_dims),
            nn.ReLU(),
            nn.Linear(fc2_dims, n_actions),
            nn.Softmax(dim=-1)
        )

        self.optimizer = optim.Adam(self.parameters(), lr=alpha)
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        self.to(self.device)

    def forward(self, state):
        dist = self.actor(state)
        dist = Categorical(dist)

        return dist

    def save_checkpoint(self):
        torch.save(self.state_dict(), self.checkpoint_file)

    def load_checkpoint(self):
        self.load_state_dict(torch.load(self.checkpoint_file))


class CriticNetwork(nn.Module):
    def __init__(self, input_dims, alpha, fc1_dims=256, fc2_dims=256,
                 chkpt_dir='tmp/ppo'):
        super(CriticNetwork, self).__init__()

        self.checkpoint_file = os.path.join(chkpt_dir, 'critic_torch_ppo')
        self.critic = nn.Sequential(
            nn.Linear(*input_dims, fc1_dims),
            nn.ReLU(),
            nn.Linear(fc1_dims, fc2_dims),
            nn.ReLU(),
            nn.Linear(fc2_dims, 1)
        )

        self.optimizer = optim.Adam(self.parameters(), lr=alpha)
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        self.to(self.device)

    def forward(self, state):
        value = self.critic(state)

        return value

    def save_checkpoint(self):
        torch.save(self.state_dict(), self.checkpoint_file)

    def load_checkpoint(self):
        self.load_state_dict(torch.load(self.checkpoint_file))


class ActorConvNetwork(nn.Module):
    def __init__(self, n_actions, input_dims, alpha,
                 fc1_dims=256, fc2_dims=256, chkpt_dir='tmp/ppo'):
        super(ActorConvNetwork, self).__init__()

        self.checkpoint_file = os.path.join(chkpt_dir, 'actor_torch_ppo')

        """
        self.conv1 = nn.Conv2d(4, 32, kernel_size=8, stride=4)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)

        self.bn = nn.BatchNorm2d(32)
        self.bn2 = nn.BatchNorm2d(64)
        self.bn3 = nn.BatchNorm2d(64)

        self.fc1 = nn.Linear(7 * 7 * 64, 512)
        self.adv = nn.Linear(512, action_size)
        """

        self.actor = nn.Sequential(
            nn.Conv2d(4, 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(7 * 7 * 64, 512),
            nn.ReLU(),
            nn.Linear(512, n_actions),
            nn.Softmax(dim=-1)
        )

        self.optimizer = optim.Adam(self.parameters(), lr=alpha)
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        self.to(self.device)

    def forward(self, state):
        dist = self.actor(state)
        dist = Categorical(dist)

        return dist

    def save_checkpoint(self):
        torch.save(self.state_dict(), self.checkpoint_file)

    def load_checkpoint(self):
        self.load_state_dict(torch.load(self.checkpoint_file))


class CriticConvNetwork(nn.Module):
    def __init__(self, input_dims, alpha, fc1_dims=256, fc2_dims=256,
                 chkpt_dir='tmp/ppo'):
        super(CriticConvNetwork, self).__init__()

        self.checkpoint_file = os.path.join(chkpt_dir, 'critic_torch_ppo')
        self.critic = nn.Sequential(
            nn.Conv2d(4, 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(7 * 7 * 64, 512),
            nn.ReLU(),
            nn.Linear(512, 1)
        )

        self.optimizer = optim.Adam(self.parameters(), lr=alpha)
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        self.to(self.device)

    def forward(self, state):
        value = self.critic(state)

        return value

    def save_checkpoint(self):
        torch.save(self.state_dict(), self.checkpoint_file)

    def load_checkpoint(self):
        self.load_state_dict(torch.load(self.checkpoint_file))


class ActorConvNetwork_v2(nn.Module):

    def __init__(self, n_actions, input_dims, alpha,
                 fc1_dims=256, fc2_dims=256, chkpt_dir='tmp/ppo'):
        super(ActorConvNetwork_v2, self).__init__()

        def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
            nn.init.orthogonal_(layer.weight, std)
            nn.init.constant_(layer.bias, bias_const)
            return layer

        self.checkpoint_file = os.path.join(chkpt_dir, 'actor_torch_ppo')

        self.actor = nn.Sequential(
            layer_init(nn.Conv2d(4, 32, kernel_size=8, stride=4)),
            nn.ReLU(),
            layer_init(nn.Conv2d(32, 64, kernel_size=4, stride=2)),
            nn.ReLU(),
            layer_init(nn.Conv2d(64, 64, kernel_size=3, stride=1)),
            nn.ReLU(),
            nn.Flatten(),
            layer_init(nn.Linear(7 * 7 * 64, 512)),
            nn.ReLU(),
            layer_init(nn.Linear(512, n_actions), std=0.01),
            nn.Softmax(dim=-1)
        )

        self.optimizer = Lion(self.parameters(), lr=alpha)
        self.scheduler = StepLR(self.optimizer, step_size=1000, gamma=0.95)
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        self.to(self.device)

    def forward(self, state):
        dist = self.actor(state).squeeze()
        dist = Categorical(dist)

        return dist

    def save_checkpoint(self):
        torch.save(self.state_dict(), self.checkpoint_file)

    def load_checkpoint(self):
        self.load_state_dict(torch.load(self.checkpoint_file))


class CriticConvNetwork_v2(nn.Module):
    def __init__(self, input_dims, alpha, fc1_dims=256, fc2_dims=256,
                 chkpt_dir='tmp/ppo'):
        super(CriticConvNetwork_v2, self).__init__()

        def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
            nn.init.orthogonal_(layer.weight, std)
            nn.init.constant_(layer.bias, bias_const)
            return layer

        self.checkpoint_file = os.path.join(chkpt_dir, 'critic_torch_ppo')
        self.critic = nn.Sequential(
            layer_init(nn.Conv2d(4, 32, kernel_size=8, stride=4)),
            nn.ReLU(),
            layer_init(nn.Conv2d(32, 64, kernel_size=4, stride=2)),
            nn.ReLU(),
            layer_init(nn.Conv2d(64, 64, kernel_size=3, stride=1)),
            nn.ReLU(),
            nn.Flatten(),
            layer_init(nn.Linear(7 * 7 * 64, 512)),
            nn.ReLU(),
            layer_init(nn.Linear(512, 1), std=1)
        )

        self.optimizer = Lion(self.parameters(), lr=alpha)
        self.scheduler = StepLR(self.optimizer, step_size=1000, gamma=0.95)
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        self.to(self.device)

    def forward(self, state):
        value = self.critic(state)

        return value

    def save_checkpoint(self):
        torch.save(self.state_dict(), self.checkpoint_file)

    def load_checkpoint(self):
        self.load_state_dict(torch.load(self.checkpoint_file))


import numpy as np
class PPOMemory:
    def __init__(self, batch_size):
        self.states = []
        self.probs = []
        self.vals = []
        self.actions = []
        self.rewards = []
        self.dones = []

        self.batch_size = batch_size

    def generate_batches(self):
        n_states = len(self.states)
        batch_start = np.arange(0, n_states, self.batch_size)
        indices = np.arange(n_states, dtype=np.int64)
        np.random.shuffle(indices)
        batches = [indices[i:i+self.batch_size] for i in batch_start]

        return np.array(self.states),\
                np.array(self.actions),\
                np.array(self.probs),\
                np.array(self.vals),\
                np.array(self.rewards),\
                np.array(self.dones),\
                batches

    def store_memory(self, states, actions, probs, vals, rewards, dones):
        self.states.extend(states)
        self.actions.extend(actions)
        self.probs.extend(probs)
        self.vals.extend(vals)
        self.rewards.extend(rewards)
        self.dones.extend(dones)

    def clear_memory(self):
        self.states = []
        self.probs = []
        self.actions = []
        self.rewards = []
        self.dones = []
        self.vals = []


total_number_steps = 1e7
learning_iteration = 0
num_envs = 8
num_step_per_iteration = num_envs * 2048
import_models = False
render = None


if __name__ == '__main__':
    def create_env(seed):
        def _init():
            env = make_atari("PongNoFrameskip-v4", 10000, render=render)
            env = wrap_deepmind(env, True, True, True, True)
            env.seed(seed)
            return env
        return _init


    envs = SyncVectorEnv([create_env(i) for i in range(num_envs)])
    agent = Agent(envs.single_action_space.n, [4, 84, 84], batch_size=64, import_models=import_models)

    print(f"Training for {total_number_steps} steps.\n"
          f"---------------------\n"
          f"Parameter Recap\n"
          f"Number of step for iteration: {num_step_per_iteration}\n"
          f"Number of concurrent environments: {num_envs}\n"
          f"Step per environment before learning iteration: {num_step_per_iteration/num_envs}\n"
          f"---------------------\n"
          f"GLHF\n"
          )

    current_rewards = np.zeros(num_envs)  # Storing Environment Rewards for logging purposes.
    cumulative_rewards = np.zeros(num_envs)  # Tracking Reward for each environment.
    done_flags = np.zeros(num_envs, dtype=bool)

    start_training_time = time.time()  # Tracking time for logging purposes.

    observations, infos = envs.reset()  # First observation, start of training for ALL environments.
    observations = np.transpose(observations, [0, 3, 1, 2])  # [num_envs, channels, height, width]
    # TODO: Write wrapper to obtain transposed observation [num_envs, num_stack, h, w]

    for step in range(0, int(total_number_steps), num_envs):

        actions, probs, vals = agent.choose_action(observations)  # Sampling Action
        observations_, rewards, terminated, truncated, infos = envs.step(actions)  # Stepping the environment

        observations_ = np.transpose(observations_, [0, 3, 1, 2])  # As Above.
        dones = np.logical_or(terminated, truncated)  # done = truncated OR terminated, gym 0.26+
        cumulative_rewards += rewards  # [num_envs] Tracking scores.

        agent.remember(observations, actions, probs, vals, rewards, dones)
        # State, action, log_probs, values, rewards, dones. [num_envs] each.

        if np.any(dones):
            indices = np.where(dones)[0]
            current_rewards[indices] = cumulative_rewards[indices]
            cumulative_rewards[indices] = 0.0
            if np.any(terminated[indices]):
                pippo = np.transpose(infos['final_observation'][indices][0], [2, 0, 1])
                rewards[indices] = agent.gamma * agent.critic(torch.tensor([pippo], dtype=torch.float32).to(agent.actor.device))[0].detach().cpu()

        if step % num_step_per_iteration == 0 and step != 0:

            print(f"Experience Collected. Learning.")
            agent.learn()
            learning_iteration += 1
            agent.save_models()
            mean_reward = np.mean(current_rewards)
            print(f"Learning Iteration: {learning_iteration}\n"
                  f"Mean Reward: {mean_reward}\n"
                  f"Step: {step}\n")

            try:
                asyncio.get_event_loop().run_until_complete(
                    send_message_to_channel(
                        f"PongNoFrameskip-v4 - MultiEnv\n"
                        f"Learning Iteration: {learning_iteration}\n"
                        f"Reward: {mean_reward}\n"
                        f"Elapsed Time {(time.time() - start_training_time)/60:.2f}\n"
                    )
                )
            except:
                pass

        # Update the observations for environments that are not done
        observations = observations_



Can you check if my PPO implementation is correct? Can you provide explanations of the problems and their solution?