improve below DAPO it is not wrong properly , only logic error , but code is runing 
import sys
import logging
from typing import Tuple, List, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import gymnasium as gym
import torch.nn.functional as F
from torch.distributions import Categorical

# Configure logging
logging.basicConfig(
    level=logging.INFO,  # Change to DEBUG for more detailed logs
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)


class PolicyNetwork(nn.Module):
    """
    Policy Network to parameterize the policy π_θ.
    """

    def __init__(self, state_dim: int, action_dim: int, hidden_dims: Tuple[int, ...] = (128, 128)) -> None:
        super(PolicyNetwork, self).__init__()
        layers: List[nn.Module] = []
        input_dim: int = state_dim
        for hidden_dim in hidden_dims:
            layers.append(nn.Linear(input_dim, hidden_dim))
            layers.append(nn.ReLU())
            input_dim = hidden_dim
        layers.append(nn.Linear(input_dim, action_dim))
        self.network: nn.Sequential = nn.Sequential(*layers)

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        logits: torch.Tensor = self.network(state)
        return logits  # Return logits; softmax applied later


class CriticNetwork(nn.Module):
    """
    Critic Network to approximate the Q-values Q(s, a).
    """

    def __init__(self, state_dim: int, action_dim: int, hidden_dims: Tuple[int, ...] = (128, 128)) -> None:
        super(CriticNetwork, self).__init__()
        layers: List[nn.Module] = []
        input_dim: int = state_dim + action_dim
        for hidden_dim in hidden_dims:
            layers.append(nn.Linear(input_dim, hidden_dim))
            layers.append(nn.ReLU())
            input_dim = hidden_dim
        layers.append(nn.Linear(input_dim, 1))
        self.network: nn.Sequential = nn.Sequential(*layers)

    def forward(self, state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
        x: torch.Tensor = torch.cat([state, action], dim=-1)
        q_value: torch.Tensor = self.network(x)
        return q_value.squeeze(-1)  # Return as (batch_size,)


class ReplayBuffer:
    """
    Replay Buffer for storing and sampling experiences.
    """

    def __init__(self, capacity: int) -> None:
        self.capacity: int = capacity
        self.buffer: List[Optional[Tuple[np.ndarray, int, float, np.ndarray, bool]]] = []
        self.position: int = 0

    def push(self, state: np.ndarray, action: int, reward: float,
             next_state: np.ndarray, done: bool) -> None:
        """Saves a transition to the buffer."""
        if len(self.buffer) < self.capacity:
            self.buffer.append(None)  # type: ignore
        self.buffer[self.position] = (state, action, reward, next_state, done)
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size: int) -> Tuple[
        np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """Samples a batch of experiences from the buffer."""
        if len(self.buffer) < batch_size:
            raise ValueError("Not enough samples in buffer to sample.")
        indices = np.random.choice(len(self.buffer), batch_size, replace=False)
        batch = [self.buffer[idx] for idx in indices]
        states, actions, rewards, next_states, dones = zip(*batch)  # type: ignore
        return (
            np.array(states),
            np.array(actions),
            np.array(rewards, dtype=np.float32),
            np.array(next_states),
            np.array(dones, dtype=np.float32)
        )

    def __len__(self) -> int:
        return len(self.buffer)


class DAPOAgent:
    """
    Dual Approximation Policy Optimization (DAPO) Agent - KL Variant.
    """

    def __init__(
        self,
        state_dim: int,
        action_dim: int,
        env: gym.Env,
        gamma: float = 0.99,
        lr_policy: float = 1e-3,
        lr_critic: float = 1e-3,
        buffer_capacity: int = 100000,
        batch_size: int = 64,
        eta: float = 0.1
    ) -> None:
        self.state_dim: int = state_dim
        self.action_dim: int = action_dim
        self.env: gym.Env = env
        self.gamma: float = gamma
        self.batch_size: int = batch_size
        self.eta: float = eta

        # Initialize device
        self.device: torch.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        logger.info(f"Using device: {self.device}")

        # Initialize networks
        self.policy_net: PolicyNetwork = PolicyNetwork(state_dim, action_dim).to(self.device)
        self.critic_net: CriticNetwork = CriticNetwork(state_dim, action_dim).to(self.device)

        # Initialize optimizers
        self.policy_optimizer: optim.Optimizer = optim.Adam(self.policy_net.parameters(), lr=lr_policy)
        self.critic_optimizer: optim.Optimizer = optim.Adam(self.critic_net.parameters(), lr=lr_critic)

        # Initialize replay buffer
        self.replay_buffer: ReplayBuffer = ReplayBuffer(capacity=buffer_capacity)

    def select_action(self, state: np.ndarray) -> int:
        """
        Select an action based on the current policy.
        """
        state_tensor: torch.Tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
        with torch.no_grad():
            logits: torch.Tensor = self.policy_net(state_tensor)
            action_probs: torch.Tensor = F.softmax(logits, dim=-1)
        dist: Categorical = Categorical(probs=action_probs)
        action: int = dist.sample().item()
        return action

    def update_critic(
        self,
        states: torch.Tensor,
        actions: torch.Tensor,
        rewards: torch.Tensor,
        next_states: torch.Tensor,
        dones: torch.Tensor
    ) -> None:
        """
        Update the critic network using TD(0) learning.
        """
        actions_one_hot = F.one_hot(actions, num_classes=self.action_dim).float()
        q_values = self.critic_net(states, actions_one_hot)  # Shape: (batch_size,)

        with torch.no_grad():
            # Next action probabilities from the policy
            next_logits = self.policy_net(next_states)
            next_action_probs = F.softmax(next_logits, dim=-1)  # Shape: (batch_size, action_dim)

            # Compute expected next Q-values
            next_actions_one_hot = F.one_hot(
                torch.arange(self.action_dim).unsqueeze(0).expand(len(next_states), -1),
                num_classes=self.action_dim
            ).float().to(self.device)  # Shape: (batch_size, action_dim, action_dim)

            next_states_expanded = next_states.unsqueeze(1).expand(-1, self.action_dim, -1).reshape(-1, self.state_dim)
            next_actions_expanded = next_actions_one_hot.reshape(-1, self.action_dim)

            next_q_values_all = self.critic_net(next_states_expanded, next_actions_expanded)  # Shape: (batch_size * action_dim,)
            next_q_values_all = next_q_values_all.view(len(next_states), self.action_dim)

            expected_next_q_values = (next_action_probs * next_q_values_all).sum(dim=1)  # Shape: (batch_size,)

            target_q_values = rewards + self.gamma * expected_next_q_values * (1 - dones)

        loss = F.mse_loss(q_values, target_q_values)

        self.critic_optimizer.zero_grad()
        loss.backward()
        self.critic_optimizer.step()

    def update_policy(self, states: torch.Tensor) -> None:
        """
        Update the policy network using the DAPO-KL update rule.
        """
        with torch.no_grad():
            # Compute current policy probabilities
            logits = self.policy_net(states)
            pi_k = F.softmax(logits, dim=-1)  # π^{(k)}_s

            # Compute Q-values for all actions
            actions_all = torch.arange(self.action_dim).unsqueeze(0).expand(len(states), -1).to(self.device)
            actions_all_one_hot = F.one_hot(actions_all, num_classes=self.action_dim).float()
            states_expanded = states.unsqueeze(1).expand(-1, self.action_dim, -1).reshape(-1, self.state_dim)
            actions_expanded = actions_all_one_hot.view(-1, self.action_dim)

            q_values_all = self.critic_net(states_expanded, actions_expanded).view(len(states), self.action_dim)  # Shape: (batch_size, action_dim)

            # Compute unnormalized target distribution
            exponent = -self.eta * q_values_all  # Shape: (batch_size, action_dim)
            unnormalized_target = pi_k * exponent.exp()

            # Normalize to get the target distribution
            target_distribution = unnormalized_target / unnormalized_target.sum(dim=-1, keepdim=True)

        # Get current policy logits
        logits = self.policy_net(states)

        # Compute the KL divergence loss
        log_probs = F.log_softmax(logits, dim=-1)
        loss = F.kl_div(log_probs, target_distribution, reduction='batchmean')

        self.policy_optimizer.zero_grad()
        loss.backward()
        self.policy_optimizer.step()

    def train(self, num_episodes: int = 1000, max_steps: int = 1000) -> None:
        """
        Train the agent.
        """
        try:
            for episode in range(1, num_episodes + 1):
                state, _ = self.env.reset()
                episode_reward = 0.0
                for step in range(max_steps):
                    action = self.select_action(state)
                    next_state, reward, terminated, truncated, _ = self.env.step(action)
                    done = terminated or truncated
                    self.replay_buffer.push(state, action, reward, next_state, done)
                    state = next_state
                    episode_reward += reward

                    if len(self.replay_buffer) > self.batch_size:
                        # Sample a batch
                        try:
                            batch = self.replay_buffer.sample(self.batch_size)
                            states_np, actions_np, rewards_np, next_states_np, dones_np = batch
                        except ValueError as e:
                            logger.warning(f"Not enough samples to train: {e}")
                            continue

                        states = torch.FloatTensor(states_np).to(self.device)
                        actions = torch.LongTensor(actions_np).to(self.device)
                        rewards = torch.FloatTensor(rewards_np).to(self.device)
                        next_states = torch.FloatTensor(next_states_np).to(self.device)
                        dones = torch.FloatTensor(dones_np).to(self.device)

                        # Update Critic
                        self.update_critic(states, actions, rewards, next_states, dones)

                        # Update Policy
                        self.update_policy(states)

                    if done:
                        break

                logger.info(f"Episode {episode} Reward: {episode_reward:.2f}")

                if episode % 50 == 0:
                    self.evaluate(num_episodes=5)

        except KeyboardInterrupt:
            logger.info("Training interrupted by user.")
        except Exception as e:
            logger.error(f"Error during training: {e}")

    def evaluate(self, num_episodes: int = 5) -> None:
        """
        Evaluate the current policy without exploration.
        """
        try:
            total_reward: float = 0.0
            for _ in range(num_episodes):
                state, _ = self.env.reset()
                done = False
                episode_reward: float = 0.0
                while not done:
                    state_tensor: torch.Tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
                    with torch.no_grad():
                        logits: torch.Tensor = self.policy_net(state_tensor)
                        action_probs: torch.Tensor = F.softmax(logits, dim=-1)
                        action: int = action_probs.argmax(dim=-1).item()
                    state, reward, terminated, truncated, _ = self.env.step(action)
                    done = terminated or truncated
                    episode_reward += reward
                total_reward += episode_reward
            average_reward = total_reward / num_episodes
            logger.info(f"Evaluation over {num_episodes} episodes: Average Reward: {average_reward:.2f}")
        except Exception as e:
            logger.error(f"Error during evaluation: {e}")


def main() -> None:
    """
    Main function to initialize environment and start training.
    """
    try:
        env_name: str = "CartPole-v1"
        env: gym.Env = gym.make(env_name)
        state_dim: int = env.observation_space.shape[0]
        action_dim: int = env.action_space.n

        agent = DAPOAgent(
            state_dim=state_dim,
            action_dim=action_dim,
            env=env,
            gamma=0.99,
            lr_policy=1e-3,
            lr_critic=1e-3,
            buffer_capacity=100000,
            batch_size=64,
            eta=0.1
        )

        agent.train(num_episodes=500, max_steps=500)
    except Exception as e:
        logger.error(f"Error in main: {e}")
    finally:
        try:
            env.close()
            logger.info("Environment closed.")
        except Exception as e:
            logger.error(f"Error closing environment: {e}")


if __name__ == "__main__":
    main()

2024-11-11 22:47:34,673 [INFO] Episode 495 Reward: 9.00
2024-11-11 22:47:34,706 [INFO] Episode 496 Reward: 8.00
2024-11-11 22:47:34,745 [INFO] Episode 497 Reward: 9.00
2024-11-11 22:47:34,778 [INFO] Episode 498 Reward: 10.00
2024-11-11 22:47:34,823 [INFO] Episode 499 Reward: 10.00
2024-11-11 22:47:34,873 [INFO] Episode 500 Reward: 11.00
2024-11-11 22:47:34,873 [INFO] Evaluation over 5 episodes: Average Reward: 9.80      average error should above 950 or  but coming very low