in the below observation space may be text, image , numerical, categories data 
action space also may be
using transformers module update (from this module any mode can act like critic , actor still our ppo should work
import copy
import math
import random
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import warnings

# Suppress specific warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

# Set a random seed for reproducibility
torch.manual_seed(42)
np.random.seed(42)
random.seed(42)


class Actor(nn.Module):
    """
    Actor network that outputs action probabilities.
    """

    def __init__(self, state_dim: int, action_dim: int, hidden_sizes: Tuple[int, ...] = (128, 128)) -> None:
        """
        Initializes the Actor network.

        Args:
            state_dim (int): Dimension of the input state.
            action_dim (int): Dimension of the output action space.
            hidden_sizes (Tuple[int, ...], optional): Sizes of hidden layers. Defaults to (128, 128).
        """
        super(Actor, self).__init__()
        layers = []
        input_dim = state_dim
        for hidden_size in hidden_sizes:
            layers.append(nn.Linear(input_dim, hidden_size))
            layers.append(nn.ReLU())
            input_dim = hidden_size
        layers.append(nn.Linear(input_dim, action_dim))
        # Removed Softmax to handle logits directly for numerical stability
        self.network = nn.Sequential(*layers)

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        """
        Forward pass through the network.

        Args:
            state (torch.Tensor): Input state.

        Returns:
            torch.Tensor: Action logits.
        """
        return self.network(state)


class Critic(nn.Module):
    """
    Critic network that estimates the state-value function.
    """

    def __init__(self, state_dim: int, hidden_sizes: Tuple[int, ...] = (128, 128)) -> None:
        """
        Initializes the Critic network.

        Args:
            state_dim (int): Dimension of the input state.
            hidden_sizes (Tuple[int, ...], optional): Sizes of hidden layers. Defaults to (128, 128).
        """
        super(Critic, self).__init__()
        layers = []
        input_dim = state_dim
        for hidden_size in hidden_sizes:
            layers.append(nn.Linear(input_dim, hidden_size))
            layers.append(nn.ReLU())
            input_dim = hidden_size
        layers.append(nn.Linear(input_dim, 1))
        self.network = nn.Sequential(*layers)

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        """
        Forward pass through the network.

        Args:
            state (torch.Tensor): Input state.

        Returns:
            torch.Tensor: Estimated state value.
        """
        return self.network(state)


@dataclass
class PPOConfig:
    """
    Configuration parameters for PPO.
    """
    state_dim: int
    action_dim: int
    hidden_sizes: Tuple[int, ...] = (128, 128)
    lr_actor: float = 3e-4
    lr_critic: float = 1e-3
    gamma: float = 0.99
    lambda_: float = 0.95
    epsilon: float = 0.2
    entropy_coef: float = 0.01
    value_coef: float = 0.5
    max_grad_norm: float = 0.5
    epochs: int = 4  # Reduced epochs for better generalization
    batch_size: int = 64
    clip_vloss: bool = True  # Clipping value loss


class PPO:
    """
    Proximal Policy Optimization (PPO) agent.
    """

    def __init__(self, config: PPOConfig) -> None:
        """
        Initializes the PPO agent with actor and critic networks.

        Args:
            config (PPOConfig): PPO configuration parameters.
        """
        self.config = config
        try:
            self.actor = Actor(config.state_dim, config.action_dim, config.hidden_sizes)
            self.critic = Critic(config.state_dim, config.hidden_sizes)
        except Exception as e:
            raise RuntimeError(f"Error initializing networks: {e}")

        self.optimizer_actor = optim.Adam(self.actor.parameters(), lr=config.lr_actor)
        self.optimizer_critic = optim.Adam(self.critic.parameters(), lr=config.lr_critic)

        self.memory: List[Dict[str, Any]] = []
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.actor.to(self.device)
        self.critic.to(self.device)

    def select_action(self, state: np.ndarray) -> Tuple[int, float, float]:
        """
        Selects an action based on the current policy.

        Args:
            state (np.ndarray): Current state.

        Returns:
            Tuple[int, float, float]: Action, log probability, and state value.
        """
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
        with torch.no_grad():
            logits = self.actor(state_tensor)
            dist = Categorical(logits=logits)
            action = dist.sample()
            log_prob = dist.log_prob(action)
            value = self.critic(state_tensor)
        return action.item(), log_prob.item(), value.item()

    def store_transition(
        self,
        state: np.ndarray,
        action: int,
        log_prob: float,
        reward: float,
        done: bool,
        value: float
    ) -> None:
        """
        Stores a transition in memory.

        Args:
            state (np.ndarray): Current state.
            action (int): Action taken.
            log_prob (float): Log probability of the action.
            reward (float): Reward received.
            done (bool): Whether the episode is done.
            value (float): Estimated state value.
        """
        self.memory.append({
            'state': state,
            'action': action,
            'log_prob': log_prob,
            'reward': reward,
            'done': done,
            'value': value
        })

    def compute_returns_and_advantages(
        self,
        last_value: float
    ) -> Tuple[List[float], List[float]]:
        """
        Computes returns and advantage estimates using GAE.

        Args:
            last_value (float): The value of the state following the last state.

        Returns:
            Tuple[List[float], List[float]]: Returns and advantages.
        """
        rewards = [t['reward'] for t in self.memory]
        dones = [t['done'] for t in self.memory]
        values = [t['value'] for t in self.memory]
        values.append(last_value)

        returns = []
        gae = 0
        for step in reversed(range(len(rewards))):
            delta = rewards[step] + self.config.gamma * values[step + 1] * (1 - dones[step]) - values[step]
            gae = delta + self.config.gamma * self.config.lambda_ * (1 - dones[step]) * gae
            returns.insert(0, gae + values[step])

        advantages = [ret - val for ret, val in zip(returns, values[:-1])]
        # Normalize advantages for better training stability
        advantages = (advantages - np.mean(advantages)) / (np.std(advantages) + 1e-8)
        return returns, advantages

    def update(self, last_value: float) -> None:
        """
        Updates the policy and value networks using the collected memory.

        Args:
            last_value (float): The value of the state following the last state.
        """
        if not self.memory:
            return

        returns, advantages = self.compute_returns_and_advantages(last_value)

        # Efficient tensor conversion
        states = torch.FloatTensor(np.array([t['state'] for t in self.memory])).to(self.device)
        actions = torch.LongTensor([t['action'] for t in self.memory]).to(self.device)
        old_log_probs = torch.FloatTensor([t['log_prob'] for t in self.memory]).to(self.device)
        returns = torch.FloatTensor(returns).to(self.device)
        advantages = torch.FloatTensor(advantages).to(self.device)

        # Create dataset and dataloader
        dataset = torch.utils.data.TensorDataset(states, actions, old_log_probs, returns, advantages)
        loader = torch.utils.data.DataLoader(dataset, batch_size=self.config.batch_size, shuffle=True)

        for _ in range(self.config.epochs):
            for batch_states, batch_actions, batch_old_log_probs, batch_returns, batch_advantages in loader:
                # Forward passes
                logits = self.actor(batch_states)
                dist = Categorical(logits=logits)
                entropy = dist.entropy().mean()
                new_log_probs = dist.log_prob(batch_actions)
                ratio = (new_log_probs - batch_old_log_probs).exp()

                # PPO Clipped Objective
                surr1 = ratio * batch_advantages
                surr2 = torch.clamp(ratio, 1.0 - self.config.epsilon, 1.0 + self.config.epsilon) * batch_advantages
                actor_loss = -torch.min(surr1, surr2).mean() - self.config.entropy_coef * entropy

                # Critic loss
                values = self.critic(batch_states).squeeze()
                critic_loss = nn.MSELoss()(values, batch_returns)
                if self.config.clip_vloss:
                    clipped_values = self.critic(batch_states).squeeze()
                    clipped_values = values + torch.clamp(clipped_values - values, -self.config.epsilon, self.config.epsilon)
                    critic_loss_1 = nn.MSELoss()(clipped_values, batch_returns)
                    critic_loss_2 = nn.MSELoss()(values, batch_returns)
                    critic_loss = torch.max(critic_loss_1, critic_loss_2)
                critic_loss = self.config.value_coef * critic_loss

                # Combined loss
                loss = actor_loss + critic_loss

                # Update actor network
                self.optimizer_actor.zero_grad()
                actor_loss.backward()
                nn.utils.clip_grad_norm_(self.actor.parameters(), self.config.max_grad_norm)
                self.optimizer_actor.step()

                # Update critic network
                self.optimizer_critic.zero_grad()
                critic_loss.backward()
                nn.utils.clip_grad_norm_(self.critic.parameters(), self.config.max_grad_norm)
                self.optimizer_critic.step()

        # Clear memory after update
        self.memory.clear()

    def save(self, path: str) -> None:
        """
        Saves the actor and critic models.

        Args:
            path (str): Path to save the models.
        """
        try:
            torch.save({
                'actor_state_dict': self.actor.state_dict(),
                'critic_state_dict': self.critic.state_dict()
            }, path)
            print(f"Model saved successfully at {path}.")
        except Exception as e:
            raise IOError(f"Error saving models: {e}")

    def load(self, path: str) -> None:
        """
        Loads the actor and critic models.

        Args:
            path (str): Path to load the models from.
        """
        try:
            checkpoint = torch.load(path, map_location=self.device)
            self.actor.load_state_dict(checkpoint['actor_state_dict'])
            self.critic.load_state_dict(checkpoint['critic_state_dict'])
            print(f"Model loaded successfully from {path}.")
        except Exception as e:
            raise IOError(f"Error loading models: {e}")


# Example usage with the CartPole-v1 environment
if __name__ == "__main__":
    import gym
    from gym.error import DependencyNotInstalled

    @dataclass
    class TrainingConfig:
        """
        Configuration parameters for training.
        """
        env_name: str = "CartPole-v1"
        total_timesteps: int = 30000  # Increased timesteps for better learning
        update_timesteps: int = 2048
        seed: int = 42
        log_interval: int = 10  # Logging every 10 episodes

    def train():
        """
        Trains the PPO agent in the specified environment.
        """
        config = TrainingConfig()
        try:
            env = gym.make(config.env_name)
        except DependencyNotInstalled as e:
            print(f"Missing dependencies for environment {config.env_name}: {e}")
            return
        except Exception as e:
            print(f"Error creating environment {config.env_name}: {e}")
            return

        # Initialize the environment seed during reset
        state, info = env.reset(seed=config.seed)
        env.action_space.seed(config.seed)

        ppo_config = PPOConfig(
            state_dim=env.observation_space.shape[0],
            action_dim=env.action_space.n,
            hidden_sizes=(128, 128),
            lr_actor=3e-4,
            lr_critic=1e-3,
            gamma=0.99,
            lambda_=0.95,
            epsilon=0.2,
            entropy_coef=0.01,
            value_coef=0.5,
            max_grad_norm=0.5,
            epochs=4,
            batch_size=64,
            clip_vloss=True
        )

        agent = PPO(ppo_config)

        timestep = 0
        episode = 0
        episode_rewards: List[float] = []
        best_reward = 0

        while timestep < config.total_timesteps:
            state, info = env.reset()
            done = False
            episode_reward = 0
            while not done:
                action, log_prob, value = agent.select_action(state)
                step_output = env.step(action)

                # Handle the updated Gym API (Gym v0.26+)
                if isinstance(step_output, tuple) and len(step_output) == 5:
                    next_state, reward, terminated, truncated, _ = step_output
                    done = terminated or truncated
                else:
                    next_state, reward, done, _ = step_output

                agent.store_transition(state, action, log_prob, reward, done, value)
                state = next_state
                timestep += 1
                episode_reward += reward

                if timestep % config.update_timesteps == 0:
                    with torch.no_grad():
                        _, _, last_value = agent.select_action(state)
                    agent.update(last_value)

                if done:
                    episode += 1
                    episode_rewards.append(episode_reward)
                    if episode % config.log_interval == 0:
                        avg_reward = np.mean(episode_rewards[-config.log_interval:])
                        print(f"Episode: {episode}, Average Reward (last {config.log_interval} episodes): {avg_reward:.2f}")
                        if avg_reward > best_reward:
                            best_reward = avg_reward
                            agent.save("ppo_cartpole_best.pth")
                            print(f"New best average reward: {best_reward:.2f}. Model saved.")
                    break

        env.close()
        agent.save("ppo_cartpole_final.pth")
        print("Training complete and models saved.")

    try:
        train()
    except Exception as e:
        print(f"An error occurred during training: {e}")