Following this instruction :
1. Strictly adhere to PEP-8 standards, ensuring code is clean, well-structured, and follows Python's style guide for readability.
2. Use appropriate modules like `typing` for type hints, ensuring code clarity, maintainability, and leveraging Python's standard libraries effectively.
3. Focus on time, space, robustness, scalability; write optimized, efficient code that can handle large datasets and scale seamlessly.
4. Implement comprehensive error handling, catching exceptions gracefully, providing meaningful error messages, and ensuring code robustness under unexpected conditions.
5. Ensure the developer agent writes advanced, maintainable Python code, balancing performance and readability while adhering to best practices.

TASK : below DAPO is not implemted correctly 

2024-11-11 22:37:42,926 [INFO] Episode 40 Reward: 10.00
2024-11-11 22:37:42,977 [INFO] Episode 41 Reward: 9.00
2024-11-11 22:37:43,028 [INFO] Episode 42 Reward: 9.00
2024-11-11 22:37:43,072 [INFO] Episode 43 Reward: 9.00
2024-11-11 22:37:43,116 [INFO] Episode 44 Reward: 8.00
2024-11-11 22:37:43,167 [INFO] Episode 45 Reward: 10.00
2024-11-11 22:37:43,208 [INFO] Episode 46 Reward: 9.00
2024-11-11 22:37:43,267 [INFO] Episode 47 Reward: 9.00
2024-11-11 22:37:43,344 [INFO] Episode 48 Reward: 10.00
2024-11-11 22:37:43,394 [INFO] Episode 49 Reward: 9.00
2024-11-11 22:37:43,442 [INFO] Episode 50 Reward: 9.00
2024-11-11 22:37:43,462 [INFO] Evaluation over 5 episodes: Average Reward: 9.40
2024-11-11 22:37:43,512 [INFO] Episode 51 Reward: 9.00
2024-11-11 22:37:43,575 [INFO] Episode 52 Reward: 10.00
2024-11-11 22:37:43,656 [INFO] Episode 53 Reward: 10.00
2024-11-11 22:37:43,709 [INFO] Episode 54 Reward: 10.00
2024-11-11 22:37:43,777 [INFO] Episode 55 Reward: 9.00
2024-11-11 22:37:43,825 [INFO] Episode 56 Reward: 10.00
2024-11-11 22:37:43,893 [INFO] Episode 57 Reward: 9.00
2024-11-11 22:37:43,977 [INFO] Episode 58 Reward: 9.00
2024-11-11 22:37:44,057 [INFO] Episode 59 Reward: 9.00
so on ....

i am giving all ideas to you inclusing sample code make its works bestly .

This section introduces the fundamental concepts of mirror descent, including Bregman divergence, Bregman projection, and the general mirror descent algorithm.

*   **Value Function (Equation 1):**
    $$
    V^{\pi}_{\rho} = E_{s \sim \rho}[V^{\pi}_s],
    $$

    where:

    *   $V^{\pi}_{\rho}$ is the value function for policy π and initial state distribution ρ
    *   $V^{\pi}_s$ is the value function for policy π starting from state s
*   **Value Function (Equation 1 continued):**

    $$
    V^{\pi}_s = E_{a_t \sim \pi_{s_t}} \left[ \sum_{t=0}^\infty \gamma^t c(s_t, a_t) | s_0 = s \right],
    $$
    where:

    *   $\gamma$ is the discount factor
    *   $c(s_t, a_t)$ is the cost at time t for state $s_t$ and action $a_t$
*   **Q-value Function (Equation 2):**

    $$
    Q^{\pi}_{s, a} = E_{a_t \sim \pi_{s_t}} \left[ \sum_{t=0}^\infty \gamma^t c(s_t, a_t) | s_0 = s, a_0 = a \right],
    $$

    where $Q^{\pi}_{s, a}$ is the Q-value function for policy π, state s, and action a.

*   **Relationship between Value and Q-value Functions:**

    $$
    V^{\pi}_s = \langle Q^{\pi}_s, \pi_s \rangle,
    $$

    where:

    *   $Q^{\pi}_s$ is the vector of Q-values for all actions in state s
    *   $\pi_s$ is the policy (probability distribution over actions) in state s

*   **Discounted State-Visitation Distribution (Equation 3):**
    $$
    d^{\pi}_{\rho, s} = (1 - \gamma) \sum_{t=0}^\infty \gamma^t P(s_t = s | s_0 \sim \rho, \pi),
    $$

    where:

    *   $d^{\pi}_{\rho, s}$ is the discounted state-visitation distribution for policy π and initial state distribution ρ
    *   $P(s_t = s | s_0 \sim \rho, \pi)$ is the probability of being in state s at time t given the initial state distribution ρ and policy π
*   **Policy Gradient (Equation 4):**

    $$
    \nabla_s V^{\pi}_{\rho} = \frac{1}{1-\gamma} d^{\pi}_{\rho, s} Q^{\pi}_s \in \mathbb{R}^{\lvert A \rvert},
    $$

    where $\nabla_s V^{\pi}_{\rho}$ is the gradient of the value function with respect to the policy in state s.

*   **Bregman Divergence (Equation 5):**

    $$
    D_{\Phi}(x, y) = \Phi(x) - \Phi(y) - \langle \nabla \Phi(y), x-y \rangle,
    $$

    where:

    *   $D_{\Phi}(x, y)$ is the Bregman divergence between x and y induced by the convex function Φ
    *   $\nabla \Phi(y)$ is the gradient of Φ at y
*   **Bregman Projection (Equation 6):**
    $$
    \text{proj}^{\mathcal{C}}_{\Phi}(y) = \arg\min_{x \in \mathcal{C}} D_{\Phi}(x, y),
    $$

    where:

    *   $\text{proj}^{\mathcal{C}}_{\Phi}(y)$ is the Bregman projection of y onto the closed convex set  $\mathcal{C}$
*   **Mirror Descent Update (Equation 7):**
    $$
    \nabla \Phi(y^{(k+1)}) = \nabla \Phi(x^{(k)}) - \eta_k g^{(k)},
    $$
    where:

    *   $x^{(k)}$ is the iterate at iteration k
    *   $y^{(k+1)}$ is the intermediate update at iteration k+1
    *   $\eta_k$ is the step size
    *   $g^{(k)}$ is the gradient of the objective function at $x^{(k)}$

*   **Mirror Descent Update (Equation 8):**

    $$
    x^{(k+1)} = \arg\min_{x \in \mathcal{C}} D_{\Phi}(x, \nabla \Phi^*(\nabla \Phi(x^{(k)}) - \eta_k g^{(k)})),
    $$

    where $\Phi^*$ is the conjugate function of Φ.

*   **Mirror Descent Update (Equation 9):**

    $$
    x^{(k+1)} = \arg\min_{x \in \mathcal{C}} \{\eta_k \langle g^{(k)}, x \rangle + D_{\Phi}(x, x^{(k)})\}.
    $$

*   **KL Divergence (Equation 10):**

    $$
    D_{\Phi}(x, y) = \sum_i (x_i \log(x_i/y_i) - x_i + y_i),
    $$

    where:

    *   x and y are probability distributions
    *   $D_{\Phi}(x, y)$ is the KL divergence between x and y when Φ is the negative entropy
*   **Mirror Descent Update with KL Divergence (Equation 11):**

    $$
    x^{(k+1)} = x^{(k)}\exp(-\eta_k g^{(k)}) / \|x^{(k)}\exp(-\eta_k g^{(k)})\|_1.
    $$
*   **KL Divergence (Equation 12):**

    $$
    D_{\Phi}(x, y) = \sum_i x_i \log(x_i/y_i),
    $$

    where x and y are probability distributions on the simplex.

*   **Mirror Descent Update with KL Divergence (Equation 13):**

    $$
    x^{(k+1)} = \nabla \Phi^*(\nabla \Phi(x^{(k)}) - \eta_k g^{(k)}).
    $$

**Section 3: Policy Optimization with Dual Function Approximation**

This section introduces DAPO, discussing how it incorporates function approximation into policy mirror descent and the importance of the dual Bregman divergence.

*   **Policy Mirror Descent Update (Equation 13):**

    $$
    \pi^{(k+1)}_s = \arg\min_{\pi_s \in \Delta(A)} \{\eta_k \langle \hat{Q}^{(k)}_s, \pi_s \rangle + D_{\Phi}(\pi_s, \pi^{(k)}_s)\}, s \in S,
    $$

    where $\hat{Q}^{(k)}_s$ is an approximation of the Q-value function at iteration k.

*   **Policy Mirror Descent Update with Function Approximation (Equation 14):**

    $$
    \theta^{(k+1)} = \arg\min_{\theta \in \Theta} E_{s \sim d^{(k)}_{\rho}} \left[E_{a \sim \pi_{\theta_s}} [\hat{Q}^{(k)}_{s, a}] + D_{\Phi}(\pi^{\theta}_s, \pi^{(k)}_s) \right],
    $$

    where:

    *   $\pi_{\theta}$ is the parameterized policy
    *   $\Theta$ is the set of policy parameters
*   **Softmax Policy Class (Equation 15):**

    $$
    \pi_{\theta_{s, a}} = \exp(f_{\theta_{s, a}}) / \|\exp(f_{\theta_s})\|_1, (s, a) \in S \times A.
    $$

*   **AMPO Approximation Loss (Equation 16):**

    $$
    \min_{\theta} E_{s \sim d^{(k)}_{\rho}} \left[ \|f_{\theta_s} - (\nabla \Phi(\pi^{(k)}_s) - \eta_k \hat{Q}^{(k)}_s) \|^2_2 \right].
    $$

*   **DAPO Approximation Loss (Equation 17):**

    $$
    \min_{\theta \in \Theta} E_{s \sim d^{(k)}_{\rho}} \left[ D_{\Phi^*}(\nabla \Phi(\pi^{(k)}_s) - \eta_k \hat{Q}^{(k)}_s, f_{\theta_s}) \right].
    $$

*   **Duality of Bregman Divergences (Equation 18):**

    $$
    D_{\Phi^*}(\nabla \Phi(\pi^{(k)}_s) - \eta_k \hat{Q}^{(k)}_s, f_{\theta_s}) = D_{\Phi}(\nabla \Phi^*(f_{\theta_s}), \nabla \Phi^*(\nabla \Phi(\pi^{(k)}_s) - \eta_k \hat{Q}^{(k)}_s)).
    $$

*   **DAPO-L2 Approximation Loss (Equation 19):**

    $$
    \min_{\theta} E_{s \sim d^{(k)}_{\rho}} [\|f_{\theta_s} - \pi^{(k)}_s + \eta_k \hat{Q}^{(k)}_s\|^2_2].
    $$

*   **DAPO-KL\* Approximation Loss (Equation 20):**
    $$
    E_{s \sim d^{(k)}_{\rho}} \left[D_{KL}(\exp(f_{\theta_s}) ‖ \pi^{(k)}_s \exp(-\eta_k\hat{Q}^{(k)}_s)) \right].
    $$

*   **DAPO-KL Approximation Loss (Equation 21):**

    $$
    E_{s \sim d^{(k)}_{\rho}} \left[D_{KL} \left(\pi_{\theta s} ‖ \pi^{(k)}_s \exp(-\eta_k \hat{Q}^{(k)}_s) / Z^{(k)}_s \right) \right].
    $$

*   **AMPO Approximation Loss (Equation 22):**

    $$
    E_{s \sim d^{(k)}_{\rho}} \left[ D_{\Phi^*_2}(\nabla \Phi_1(\pi^{(k)}_s) - \eta_k \hat{Q}^{(k)}_s, f_{\theta_s}) \right],
    $$

    where:

    *   $\Phi_1$ is the negative entropy on $\mathbb{R}^{\lvert A \rvert}_+$
    *   $\Phi_2$ is the squared L2-norm.

**Section 4: Convergence Analysis**

This section presents the convergence analysis for DAPO-KL and SAC, including assumptions and the main convergence theorems.

**Section 4.3: Connection to SAC**

This section establishes the connection between DAPO-KL and the Soft Actor-Critic (SAC) algorithm, demonstrating that SAC can be viewed as a special case of DAPO-KL.

*   **Entropy-Regularized Value Function:**

    $$
    V^{\pi}_{\tau, \rho} = E_{a_t \sim \pi_{s_t}} \left[\sum_{t=0}^\infty \gamma^t (c(s_t, a_t) + \tau \log \pi(a_t | s_t)) | s_0 \sim \rho \right],
    $$

    where τ is the regularization parameter.

*   **Entropy-Regularized Q-value Function (Equation 23):**

    $$
    Q^{\pi}_{\tau}(s, a) = E_{a_t \sim \pi_{s_t}} \left[ \sum_{t=0}^\infty \gamma^t (c(s_t, a_t) + \tau \log \pi(a_t | s_t)) | s_0 = s, a_0 = a \right].
    $$
*   **Policy Gradient for Entropy-Regularized Value Function:**

    $$
    \nabla_s V^{\pi}_{\tau, \rho} = \frac{1}{1 - \gamma}d^{\pi}_{\rho, s}Q^{\pi}_{\tau, s} \in \mathbb{R}^{\lvert A \rvert}.
    $$

*   **DAPO-KL Update Rule for Entropy-Regularized MDP (Equation 24):**
    $$
    \theta^{(k+1)} \in \arg\min_{\theta} E_{s \sim d^{(k)}_{\rho}} \left[D_{KL} \left( \pi_{\theta_s} ‖ \pi^{(k)}_s \exp(-\eta_k Q^{(k)}_{\tau, s})/Z^{(k)}_s \right) \right],
    $$

    where $Z^{(k)}_s$ is the normalization factor.

*   **SAC Q-value Function (Equation 25):**
    $$
    q^{\pi}_{\tau}(s, a) = c(s, a) + \gamma E_{s' \sim P(s, a)} [V^{\pi}_{\tau}(s')],
    $$

    $$
    V^{\pi}_{\tau}(s) = E_{a' \sim \pi_s} [\tau \log \pi(a' | s) + q^{\pi}_{\tau}(s, a')].
    $$
*   **Relationship between SAC and Entropy-Regularized Q-value Functions:**

    $$
    q^{\pi}_{\tau}(s, a) = Q^{\pi}_{\tau}(s, a) - \tau \log \pi(a | s).
    $$

*   **SAC Policy Update Rule:**
    $$
    \theta^{(k+1)} \in \arg\min_{\theta} E_{s \sim d^{(k)}_{\rho}} \left[ D_{KL} \left(\pi_{\theta_s} ‖ \exp(-q^{(k)}_{\tau, s} / \tau)/Z^{(k)}_s \right) \right]
    $$

*   **SAC Policy Update Rule (continued):**

    $$
    = \arg\min_{\theta} E_{s \sim d^{(k)}_{\rho}} \left[ D_{KL} \left(\pi_{\theta_s} ‖ \exp(-Q^{(k)}_{\tau, s} / \tau + \log \pi^{(k)}_s)/Z^{(k)}_s \right) \right]
    $$

*   **SAC Policy Update Rule (continued):**

    $$
    = \arg\min_{\theta} E_{s \sim d^{(k)}_{\rho}} \left[ D_{KL} \left( \pi_{\theta_s} ‖ \pi^{(k)}_s \exp(-Q^{(k)}_{\tau, s}/\tau)/Z^{(k)}_s \right) \right].
    $$
import logging
import sys
from typing import Any, Optional, Tuple, List

import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
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 = 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(*layers)

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        logits = self.network(state)
        action_probs = F.softmax(logits, dim=-1)
        return action_probs


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 = 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(*layers)

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


class FunctionApproximator(nn.Module):
    """
    Function Approximator for dual updates f_θ(s).
    Outputs a vector of size equal to action_dim per state.
    """

    def __init__(self, state_dim: int, action_dim: int, hidden_dims: Tuple[int, ...] = (128, 128)) -> None:
        super(FunctionApproximator, self).__init__()
        layers: List[nn.Module] = []
        input_dim = 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(*layers)

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        return self.network(state)


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:
        if len(self.buffer) < self.capacity:
            self.buffer.append(None)
        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]:
        if len(self.buffer) < batch_size:
            raise ValueError("Not enough samples in buffer to sample.")
        batch_indices = np.random.choice(len(self.buffer), batch_size, replace=False)
        batch = [self.buffer[idx] for idx in batch_indices]
        states, actions, rewards, next_states, dones = zip(*batch)
        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 - Kullback-Leibler (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,
        lr_approximator: float = 1e-3,
        buffer_capacity: int = 100000,
        batch_size: int = 64,
        eta: float = 0.1,
        mirror_map: str = 'kl'
    ) -> 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
        self.mirror_map: str = mirror_map

        # Initialize device first to ensure it's available for model initialization
        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)
        self.approximator: FunctionApproximator = FunctionApproximator(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)
        self.approximator_optimizer: optim.Optimizer = optim.Adam(self.approximator.parameters(), lr=lr_approximator)

        # 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():
            action_probs: torch.Tensor = self.policy_net(state_tensor)
        dist = Categorical(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 by minimizing the MSE loss.
        """
        try:
            actions_one_hot = F.one_hot(actions, num_classes=self.action_dim).float().to(self.device)
            q_values = self.critic_net(states, actions_one_hot)

            with torch.no_grad():
                next_action_probs = self.policy_net(next_states)
                next_dist = Categorical(next_action_probs)
                next_actions = next_dist.sample()
                next_actions_one_hot = F.one_hot(next_actions, num_classes=self.action_dim).float().to(self.device)
                target_q = rewards + self.gamma * self.critic_net(next_states, next_actions_one_hot) * (1 - dones)

            loss = F.mse_loss(q_values, target_q)
            self.critic_optimizer.zero_grad()
            loss.backward()
            self.critic_optimizer.step()
            logger.debug(f"Critic Loss: {loss.item():.4f}")
        except Exception as e:
            logger.error(f"Error in critic update: {e}")

    def update_approximator(
        self,
        states: torch.Tensor,
        dual_targets: torch.Tensor
    ) -> None:
        """
        Update the function approximator f_θ(s).
        """
        try:
            predictions = self.approximator(states)
            loss = F.mse_loss(predictions, dual_targets)
            self.approximator_optimizer.zero_grad()
            loss.backward()
            self.approximator_optimizer.step()
            logger.debug(f"Approximator Loss: {loss.item():.4f}")
        except Exception as e:
            logger.error(f"Error in approximator update: {e}")

    def policy_update(
        self,
        states: torch.Tensor,
        dual_variables: torch.Tensor
    ) -> None:
        """
        Update the policy network using Mirror Descent with KL divergence.
        """
        try:
            action_probs = self.policy_net(states)

            # Compute the KL divergence between current policy and the desired distribution
            desired_dist = F.softmax(dual_variables, dim=-1)
            kl_div = F.kl_div(action_probs.log(), desired_dist, reduction='batchmean')

            loss = kl_div

            self.policy_optimizer.zero_grad()
            loss.backward()
            self.policy_optimizer.step()
            logger.debug(f"Policy Loss (KL Divergence): {loss.item():.4f}")
        except Exception as e:
            logger.error(f"Error in policy update: {e}")

    def train(self, num_episodes: int = 1000, max_steps: int = 1000) -> None:
        """
        Train the DAPO agent.
        """
        try:
            for episode in range(1, num_episodes + 1):
                state, _ = self.env.reset()
                episode_reward: float = 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:
                        continue

                    # Sample a batch
                    try:
                        states_np, actions_np, rewards_np, next_states_np, dones_np = self.replay_buffer.sample(self.batch_size)
                    except ValueError as e:
                        logger.warning(f"Skip update due to: {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)

                    # Generate dual targets
                    with torch.no_grad():
                        # Get current policy probabilities
                        pi_k = self.policy_net(states)  # Shape: (batch_size, action_dim)
                        log_pi_k = torch.log(pi_k + 1e-8)  # Add epsilon to prevent log(0)

                        # Get Q-values for all actions
                        actions_all = torch.arange(self.action_dim).unsqueeze(0).repeat(self.batch_size, 1).to(self.device)  # Shape: (batch_size, action_dim)
                        actions_all_one_hot = F.one_hot(actions_all, num_classes=self.action_dim).float()  # Shape: (batch_size, action_dim, action_dim)
                        states_expanded = states.unsqueeze(1).repeat(1, self.action_dim, 1).view(-1, self.state_dim)  # Shape: (batch_size * action_dim, state_dim)
                        actions_all_one_hot = actions_all_one_hot.view(-1, self.action_dim)  # Shape: (batch_size * action_dim, action_dim)

                        q_vals_all = self.critic_net(states_expanded, actions_all_one_hot).view(self.batch_size, self.action_dim)  # Shape: (batch_size, action_dim)

                        # Compute dual targets: f_theta = log(pi_k) - eta * Q
                        dual_targets = log_pi_k - self.eta * q_vals_all  # Shape: (batch_size, action_dim)

                    # Update Approximator
                    self.update_approximator(states, dual_targets)

                    # Update Policy
                    # Get approximated dual variables from the approximator
                    with torch.no_grad():
                        f_theta = self.approximator(states)  # Shape: (batch_size, action_dim)

                    self.policy_update(states, f_theta)

                    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 nabla_phi_star(self, q_values: torch.Tensor) -> torch.Tensor:
        """
        Compute the gradient of the conjugate mirror map Φ*.
        For KL divergence, Φ(x) = sum_i x_i log x_i - x_i,
        so Φ*(y) = log(sum_i exp(y_i)) which leads to ∇Φ*(y) = softmax(y).
        """
        try:
            return F.softmax(self.eta * q_values, dim=-1)
        except Exception as e:
            logger.error(f"Error in nabla_phi_star: {e}")
            return torch.zeros_like(q_values)

    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
                while not done:
                    state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
                    with torch.no_grad():
                        action_probs = self.policy_net(state_tensor)
                    action = torch.argmax(action_probs, dim=-1).item()
                    state, reward, terminated, truncated, _ = self.env.step(action)
                    done = terminated or truncated
                    total_reward += 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"  # Changed to CartPole-v1 as per user requirement
        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,
            lr_approximator=1e-3,
            buffer_capacity=100000,
            batch_size=64,
            eta=0.1,
            mirror_map='kl'  # Using KL divergence as Bregman divergence
        )

        agent.train(num_episodes=1000, max_steps=1000)
    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()