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 : basic code is given , we have develop correct code with correct logic ----

### Bregman Divergence and Bregman Projection


Given a convex function \(\Phi\) of Legendre type, the Bregman divergence between two points \(x\) and \(y\) (where \(x \in \text{dom}(\Phi)\), and \(y \in \text{int(dom}(\Phi))\)) is defined as:
\[
D_\Phi(x, y) = \Phi(x) - \Phi(y) - \langle \nabla \Phi(y), x - y \rangle.
\]
This measures the difference between \(\Phi(x)\) and the first-order approximation of \(\Phi(x)\) around \(y\), capturing how different the structure of \(x\) and \(y\) is with respect to \(\Phi\)'s curvature.

The **Bregman projection** projects a point \(y\) onto a convex set \(C\) under the Bregman divergence induced by \(\Phi\):
\[
\text{proj}_{\Phi, C}(y) = \arg \min_{x \in C} D_\Phi(x, y).
\]
This operation is crucial in the mirror descent update step when we need to ensure iterates remain within a feasible region \(C\).

### Mirror Descent Algorithm

Given a convex function \(f\) and a closed convex set \(C\), MD proceeds in two main steps at each iteration \(k\):

1. **Dual variable update**:
   \[
   \nabla \Phi(y^{(k+1)}) = \nabla \Phi(x^{(k)}) - \eta_k g^{(k)},
   \]
   where \(\eta_k\) is the step size and \( g^{(k)} \) is the gradient (or sub-gradient) of \(f\) at \(x^{(k)}\).

2. **Bregman projection**:
   \[
   x^{(k+1)} = \text{proj}_{\Phi, C}(y^{(k+1)}).
   \]

In this process, \(\Phi\) is the **mirror map**, and \(\Phi^*\) is its **conjugate function**, defined as:
\[
\Phi^*(x^*) = \sup_{x \in \text{dom}(\Phi)} \{ \langle x, x^* \rangle - \Phi(x) \},
\]
which enables the primal-dual mapping structure. Additionally, if \(\Phi\) is differentiable, the relations \(\nabla \Phi^*(\nabla \Phi(x)) = x\) hold, facilitating practical implementations of the MD algorithm.

Using Bregman divergence, the MD update can be rewritten as:
\[
x^{(k+1)} = \arg \min_{x \in C} \left\{ \eta_k \langle g^{(k)}, x \rangle + D_\Phi(x, x^{(k)}) \right\}.
\]
This highlights the two parts of the update: minimizing the linearized form of the objective function \(\langle g^{(k)}, x \rangle\) and regularizing the move by penalizing large changes in terms of the Bregman divergence.

### Examples of Mirror Descent

The text gives three specific examples of MD algorithms applied to the problem of minimizing a convex function \(f(x)\) over the simplex \(\Delta\).

#### Example 2.1: Squared L2-norm

Here, the mirror map \(\Phi(x) = \frac{1}{2} \|x\|_2^2\) corresponds to standard Euclidean projection, leading to the **projected gradient descent algorithm**:
\[
x^{(k+1)} = \arg \min_{x \in \Delta} \|x - (x^{(k)} - \eta_k \nabla f(x^{(k)}))\|_2^2,
\]
which projects the new point \(x^{(k)} - \eta_k \nabla f(x^{(k)})\) back onto the simplex \(\Delta\).

#### Example 2.2: Negative Entropy on \( \mathbb{R}_+^n \)

Using the negative entropy function \(\Phi(x) = \sum_i (x_i \log x_i - x_i)\) as the mirror map leads to an **exponentiated gradient update**:
\[
x^{(k+1)} = \frac{x^{(k)} \exp(-\eta_k g^{(k)})}{\|x^{(k)} \exp(-\eta_k g^{(k)})\|_1},
\]
where the update is performed in the dual space using a multiplicative adjustment, followed by normalization to ensure the result remains on the simplex.

#### Example 2.3: Negative Entropy on the Simplex

This example refines Example 2.2 by restricting \(\Phi(x)\) to the simplex \(\Delta\). The resulting update is similar to Example 2.2:
\[
x^{(k+1)} = \frac{x^{(k)} \exp(-\eta_k g^{(k)})}{\|x^{(k)} \exp(-\eta_k g^{(k)})\|_1},
\]
but no projection step is necessary because the mirror map naturally keeps the iterates within the simplex.

### Remarks and Differences Between Examples

- **Example 2.1** corresponds to standard Euclidean geometry and results in classical projected gradient descent.
- **Examples 2.2 and 2.3** both lead to the same exponentiated gradient update rule, though they differ in terms of the underlying structure. Example 2.2 uses the negative entropy on \( \mathbb{R}_+^n \), while Example 2.3 focuses on the simplex itself. In Example 2.3, the mirror map is not differentiable, and the standard convex optimization machinery cannot be directly applied since \(\Phi\) in this case lacks a fully invertible gradient.


## End-to-End Pseudo Algorithm Development for DAPO

The sources present DAPO (Dual Approximation Policy Optimization), a reinforcement learning framework that extends policy mirror descent methods to incorporate general function approximation. The algorithm leverages the concept of a Bregman projected policy class and utilizes the dual Bregman divergence for measuring the accuracy of function approximation. This response outlines the end-to-end pseudo-algorithm for DAPO, focusing on the DAPO-KL variant, along with its mathematical equations.

### Algorithm 1: Dual Approximation Policy Optimization (DAPO)

**Input:**

*   Initial policy π<sup>(0)</sup> with parameters θ<sup>(0)</sup>
*   Mirror map Φ
*   Learning rate schedule {η<sub>k</sub>}
*   Initial state distribution ρ ∈ Δ(S) (optional, for online training)
*   Reply buffer distribution ν<sup>(k)</sup> ∈ Δ(S) (optional, for off-policy training)

**Procedure:**

1.  **Initialization:** Set k = 0.
2.  **Loop:** Repeat until convergence or a predefined number of iterations:

    *   **Critic Update:** Find $\hat{Q}^{(k)}$ that approximates Q<sup>(k)</sup>.
    *   **Policy Update:**
        *   Find θ<sup>(k+1)</sup> by (approximately) solving:
            $$
            \min_{\theta \in \Theta} E_{s \sim \nu^{(k)}} \left[D_{\Phi^*} \left( \nabla \Phi(\pi^{(k)}_s) - \eta_k \hat{Q}^{(k)}_s, f_{\theta s} \right) \right],
            $$

            where:
            *   ν<sup>(k)</sup> can be replaced by d<sup>(k)</sup><sub>ρ</sub> for online training.
            *   $f_{\theta s}$ is a parameterized function used for approximating the dual update.

        *   Assign π<sup>(k+1)</sup><sub>s</sub> using the Bregman projection:
            $$
            \pi^{(k+1)}_s = \text{proj}^{\Delta(A)}_{\Phi} (\nabla \Phi^*(f_{\theta^{(k+1)}_s})), s \in S.
            $$
    *   Increment k: k ← k + 1.

**Output:**

*   Learned policy π<sup>(K)</sup> after K iterations

### DAPO-KL Instantiation

The DAPO-KL variant utilizes the negative entropy restricted on Δ(A) as the mirror map Φ. This choice leads to a simplified algorithm due to the range of $\nabla \Phi^*$ being Δ(A), eliminating the need for explicit projection.

**DAPO-KL Algorithm:**

1.  **Initialization:** Set k = 0.

2.  **Loop:** Repeat until convergence:

    *   **Critic Update:** Find $\hat{Q}^{(k)}$ that approximates Q<sup>(k)</sup>.
    *   **Policy Update:** Find θ<sup>(k+1)</sup> by (approximately) solving:

        $$
        \min_{\theta \in \Theta} E_{s \sim \nu^{(k)}} \left[D_{KL} \left( \pi_{\theta s} ‖ \pi^{(k)}_s \exp(-\eta_k \hat{Q}^{(k)}_s)/ Z^{(k)}_s \right) \right], 
        $$

        where:

        *   $D_{KL}$ is the KL divergence.
        *   $Z^{(k)}_s = ‖ \pi^{(k)}_s \exp(-\eta_k\hat{Q}^{(k)}_s) ‖_1$.
        *   π<sub>θs</sub> is the parameterized policy for state s.
        *   ν<sup>(k)</sup> can be replaced by d<sup>(k)</sup><sub>ρ</sub> for online training.

    *   Update the policy:
        $$
        \pi^{(k+1)}_s = \nabla \Phi^*(f^{(k+1)}_s) =  \exp(f_{\theta s, a}) / ‖ \exp(f_{\theta s}) ‖_1.
        $$
    *   Increment k: k ← k + 1.

**Output:**

*   Learned policy π<sup>(K)</sup> after K iterations

### Remarks:

*   The choice of mirror map significantly affects the form of the approximation loss and the policy update rule.
*   The critic update step involves approximating the Q-value function, a well-studied problem in reinforcement learning.
*   The learning rate schedule {η<sub>k</sub>} plays a crucial role in the convergence of the algorithm.
*   The use of reply buffer distribution ν<sup>(k)</sup> allows for off-policy training, where data is collected from a different policy than the one being optimized.
*   The convergence analysis of DAPO-KL provides linear and sublinear convergence rates depending on the learning rate schedule.

The provided algorithm and equations offer a detailed understanding of DAPO's development, highlighting its theoretical and practical aspects. The framework's flexibility in choosing mirror maps and incorporating function approximation opens up exciting avenues for designing and analyzing new policy optimization algorithms for diverse reinforcement learning applications.

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
from typing import Optional, Tuple
import logging
import sys


# Configure logging
logging.basicConfig(
    level=logging.INFO,  # Set to DEBUG for 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_dim: Tuple[int, ...] = (128, 128)) -> None:
        super(PolicyNetwork, self).__init__()
        layers = []
        input_dim = state_dim
        for hidden in hidden_dim:
            layers.append(nn.Linear(input_dim, hidden))
            layers.append(nn.ReLU())
            input_dim = hidden
        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_dim: Tuple[int, ...] = (128, 128)) -> None:
        super(CriticNetwork, self).__init__()
        layers = []
        input_dim = state_dim + action_dim
        for hidden in hidden_dim:
            layers.append(nn.Linear(input_dim, hidden))
            layers.append(nn.ReLU())
            input_dim = hidden
        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


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_dim: Tuple[int, ...] = (128, 128)) -> None:
        super(FunctionApproximator, self).__init__()
        layers = []
        input_dim = state_dim
        for hidden in hidden_dim:
            layers.append(nn.Linear(input_dim, hidden))
            layers.append(nn.ReLU())
            input_dim = hidden
        layers.append(nn.Linear(input_dim, action_dim))  # Changed from 1 to 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 = capacity
        self.buffer = []
        self.position = 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:
        batch = np.random.choice(len(self.buffer), batch_size, replace=False)
        states, actions, rewards, next_states, dones = zip(*(self.buffer[idx] for idx in batch))
        return (np.array(states), np.array(actions),
                np.array(rewards), np.array(next_states), np.array(dones))

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


class DAPOAgent:
    """
    Dual Approximation Policy Optimization (DAPO) Agent.
    """
    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 = state_dim
        self.action_dim = action_dim
        self.env = env
        self.gamma = gamma
        self.batch_size = batch_size
        self.eta = eta
        self.mirror_map = mirror_map

        # Initialize networks
        self.policy_net = PolicyNetwork(state_dim, action_dim)
        self.critic_net = CriticNetwork(state_dim, action_dim)
        self.approximator = FunctionApproximator(state_dim, action_dim)  # Pass action_dim

        # Initialize optimizers
        self.policy_optimizer = optim.Adam(self.policy_net.parameters(), lr=lr_policy)
        self.critic_optimizer = optim.Adam(self.critic_net.parameters(), lr=lr_critic)
        self.approximator_optimizer = optim.Adam(self.approximator.parameters(), lr=lr_approximator)

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

        # Initialize device
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.policy_net.to(self.device)
        self.critic_net.to(self.device)
        self.approximator.to(self.device)

    def select_action(self, state: np.ndarray) -> int:
        """
        Select an action based on the current policy.
        """
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
        with torch.no_grad():
            action_probs = self.policy_net(state_tensor)
        dist = Categorical(action_probs)
        action = 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()
            q_values = self.critic_net(states, actions_one_hot).squeeze(-1)
            with torch.no_grad():
                next_action_probs = self.policy_net(next_states)
                next_actions = Categorical(next_action_probs).sample()
                next_actions_one_hot = F.one_hot(next_actions, num_classes=self.action_dim).float()
                target_q = rewards + self.gamma * self.critic_net(next_states, next_actions_one_hot).squeeze(-1) * (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()}")
        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()}")
        except Exception as e:
            logger.error(f"Error in approximator update: {e}")

    def policy_update(self, states: torch.Tensor, q_values: torch.Tensor) -> None:
        """
        Update the policy network using Mirror Descent with KL divergence.
        """
        try:
            # Compute dual variable
            f_theta = self.approximator(states)
            dual_variable = torch.clamp(torch.log_softmax(f_theta, dim=-1), min=-20, max=20)

            # Compute mirror descent update
            action_probs = self.policy_net(states)
            # Expand q_values to match action_dim for broadcasting
            q_values_expanded = q_values.detach().unsqueeze(-1).expand(-1, self.action_dim)
            # Compute the first term: sum(action_probs * (q_values * eta))
            term1 = torch.sum(action_probs * (q_values_expanded * self.eta), dim=-1).mean()
            # Compute the KL divergence between the policy and dual_variable
            term2 = F.kl_div(action_probs.log(), torch.softmax(dual_variable, dim=-1), reduction='batchmean')
            loss = term1 + term2

            self.policy_optimizer.zero_grad()
            loss.backward()
            self.policy_optimizer.step()
            logger.debug(f"Policy Loss: {loss.item()}")
        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 = 0
                for step in range(max_steps):
                    action = self.select_action(state)
                    next_state, reward, done, truncated, _ = self.env.step(action)
                    self.replay_buffer.push(state, action, reward, next_state, done or truncated)
                    state = next_state
                    episode_reward += reward

                    if len(self.replay_buffer) < self.batch_size:
                        continue

                    # Sample a batch
                    batch = self.replay_buffer.sample(self.batch_size)
                    states, actions, rewards, next_states, dones = batch
                    states = torch.FloatTensor(states).to(self.device)
                    actions = torch.LongTensor(actions).to(self.device)
                    rewards = torch.FloatTensor(rewards).to(self.device)
                    next_states = torch.FloatTensor(next_states).to(self.device)
                    dones = torch.FloatTensor(dones).to(self.device)

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

                    # Update Approximator
                    # To align dual_targets with action_dim, compute f_theta for all actions
                    # Here, dual_targets should also have shape (batch_size, action_dim)
                    # Assuming q_vals corresponds to Q(s, a) for applicable actions
                    # However, since critic estimates Q(s,a) for specific actions, not all, this needs careful handling
                    # One approach is to compute q_vals for all possible actions
                    with torch.no_grad():
                        # Get all possible actions
                        actions_all = torch.arange(self.action_dim).unsqueeze(0).repeat(self.batch_size, 1).to(self.device)
                        actions_all_one_hot = F.one_hot(actions_all, num_classes=self.action_dim).float()
                        # Compute Q(s, a) for all actions
                        q_vals_all = self.critic_net(states.unsqueeze(1).repeat(1, self.action_dim, 1).view(-1, self.state_dim),
                                                     actions_all_one_hot.view(-1, self.action_dim)).squeeze(-1)
                        q_vals_all = q_vals_all.view(self.batch_size, self.action_dim)
                        # Compute dual_targets
                        dual_targets = self.nabla_phi_star(q_vals_all)
                    
                    self.update_approximator(states, dual_targets)

                    # Update Policy
                    self.policy_update(states, q_vals_all.mean(dim=-1))  # Using mean Q for policy update

                    if done:
                        break

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

                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).
        However, to adhere to the nature of function approximator and policy update,
        we can use exp(y) normalized appropriately.
        """
        try:
            # Apply exponential and normalize to ensure proper scaling
            exp_q = torch.exp(q_values)
            sum_exp_q = torch.sum(exp_q, dim=-1, keepdim=True)
            return exp_q / sum_exp_q
        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 = 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, done, truncated, _ = self.env.step(action)
                    total_reward += reward
            average_reward = total_reward / num_episodes
            logger.info(f"Evaluation over {num_episodes} episodes: Average Reward: {average_reward}")
        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 = "CartPole-v1"  # Changed to CartPole-v1 as per user modification
        env = gym.make(env_name, render_mode="human")
        state_dim = env.observation_space.shape[0]
        action_dim = 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:
        env.close()
        logger.info("Environment closed.")


if __name__ == "__main__":
    main()