n.exe e:/LLMS/Fine-tuning/DAPO/Mirror_Descent.py
INFO:__main__:Running example with Squared L2-norm mirror map.
INFO:__main__:Starting Mirror Descent optimization.
Traceback (most recent call last):
  File "e:\LLMS\Fine-tuning\DAPO\Mirror_Descent.py", line 414, in <module>
    main()
  File "e:\LLMS\Fine-tuning\DAPO\Mirror_Descent.py", line 409, in main
    example_squared_l2_norm()
  File "e:\LLMS\Fine-tuning\DAPO\Mirror_Descent.py", line 363, in example_squared_l2_norm
    x_opt = optimizer.optimize(x_init, objective, num_iterations, verbose=True)
  File "e:\LLMS\Fine-tuning\DAPO\Mirror_Descent.py", line 273, in optimize
    obj.backward()
  File "C:\Users\hemanthk.LAP53-FJS.000\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\_tensor.py", line 525, in backward
    torch.autograd.backward(
  File "C:\Users\hemanthk.LAP53-FJS.000\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\autograd\__init__.py", line 267, in backward
    _engine_run_backward(
  File "C:\Users\hemanthk.LAP53-FJS.000\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\autograd\graph.py", line 744, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
import torch
from torch import Tensor
from typing import Callable, Optional, Tuple, Union
import abc
import logging


# Configure logging for debug purposes
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class MirrorMap(abc.ABC):
    """
    Abstract base class for defining mirror maps.
    """

    @abc.abstractmethod
    def phi(self, x: Tensor) -> Tensor:
        """
        Compute the convex function Φ(x).

        Args:
            x (Tensor): Input tensor.

        Returns:
            Tensor: Φ(x).
        """
        pass

    @abc.abstractmethod
    def phi_star(self, x_star: Tensor) -> Tensor:
        """
        Compute the conjugate function Φ∗(x∗).

        Args:
            x_star (Tensor): Input tensor in the dual space.

        Returns:
            Tensor: Φ∗(x∗).
        """
        pass

    @abc.abstractmethod
    def grad_phi(self, x: Tensor) -> Tensor:
        """
        Compute the gradient of Φ at x, ∇Φ(x).

        Args:
            x (Tensor): Input tensor.

        Returns:
            Tensor: Gradient tensor.
        """
        pass

    @abc.abstractmethod
    def grad_phi_star(self, x_star: Tensor) -> Tensor:
        """
        Compute the gradient of Φ∗ at x∗, ∇Φ∗(x∗).

        Args:
            x_star (Tensor): Input tensor in the dual space.

        Returns:
            Tensor: Gradient tensor.
        """
        pass

    @abc.abstractmethod
    def bregman_divergence(self, x: Tensor, y: Tensor) -> Tensor:
        """
        Compute the Bregman divergence DΦ(x, y).

        Args:
            x (Tensor): Primal space tensor.
            y (Tensor): Primal space tensor.

        Returns:
            Tensor: Bregman divergence.
        """
        pass

    @abc.abstractmethod
    def projection(self, y: Tensor, constraints: Optional[Callable[[Tensor], bool]] = None) -> Tensor:
        """
        Compute the Bregman projection of y onto the convex set C.

        Args:
            y (Tensor): Input tensor in the primal space.
            constraints (Optional[Callable[[Tensor], bool]]): Constraints defining the convex set C.

        Returns:
            Tensor: Projected tensor x ∈ C.
        """
        pass


class SquaredL2NormMirrorMap(MirrorMap):
    """
    Mirror map using the squared L2-norm: Φ(x) = 0.5 * ||x||^2.
    """

    def phi(self, x: Tensor) -> Tensor:
        return 0.5 * torch.sum(x ** 2)

    def phi_star(self, x_star: Tensor) -> Tensor:
        return 0.5 * torch.sum(x_star ** 2)

    def grad_phi(self, x: Tensor) -> Tensor:
        return x

    def grad_phi_star(self, x_star: Tensor) -> Tensor:
        return x_star

    def bregman_divergence(self, x: Tensor, y: Tensor) -> Tensor:
        return 0.5 * torch.sum((x - y) ** 2)

    def projection(self, y: Tensor, constraints: Optional[Callable[[Tensor], bool]] = None) -> Tensor:
        # For Euclidean projection, PyTorch does not have a built-in projection onto generic convex sets.
        # Therefore, the user must provide the projection logic via constraints.
        if constraints is None:
            raise NotImplementedError("Projection method requires constraints.")
        # Assuming constraints define a projection function
        return constraints(y)


class NegativeEntropyMirrorMap(MirrorMap):
    """
    Mirror map using the negative entropy: Φ(x) = Σ x_i log(x_i) - x_i.
    Applicable for x ∈ R^n_+.
    """

    def phi(self, x: Tensor) -> Tensor:
        # To ensure numerical stability, add a small epsilon inside the log.
        epsilon = 1e-12
        return torch.sum(x * torch.log(x + epsilon) - x)

    def phi_star(self, x_star: Tensor) -> Tensor:
        # Φ∗(x∗) = log(Σ exp(x∗_i))
        return torch.log(torch.sum(torch.exp(x_star)))

    def grad_phi(self, x: Tensor) -> Tensor:
        # ∇Φ(x) = log(x)
        epsilon = 1e-12
        return torch.log(x + epsilon)

    def grad_phi_star(self, x_star: Tensor) -> Tensor:
        # ∇Φ∗(x∗) = exp(x∗) / ||exp(x∗)||_1
        exp_x_star = torch.exp(x_star)
        sum_exp = torch.sum(exp_x_star)
        return exp_x_star / sum_exp

    def bregman_divergence(self, x: Tensor, y: Tensor) -> Tensor:
        # KL divergence: DΦ(x, y) = Σ x_i log(x_i / y_i) - x_i + y_i
        epsilon = 1e-12
        return torch.sum(x * torch.log((x + epsilon) / (y + epsilon)) - x + y)

    def projection(self, y: Tensor, constraints: Optional[Callable[[Tensor], bool]] = None) -> Tensor:
        # Project onto the simplex: x = y / ||y||_1
        sum_y = torch.sum(y)
        if sum_y == 0:
            raise ValueError("Sum of y is zero, cannot project onto simplex.")
        return y / sum_y


class MirrorDescent:
    """
    Mirror Descent optimization algorithm.
    """

    def __init__(
        self,
        mirror_map: MirrorMap,
        step_size: Union[float, Callable[[int], float]],
        constraints: Optional[Callable[[Tensor], bool]] = None,
    ) -> None:
        """
        Initialize the Mirror Descent optimizer.

        Args:
            mirror_map (MirrorMap): The mirror map to use.
            step_size (Union[float, Callable[[int], float]]): Step size (constant or function of iteration).
            constraints (Optional[Callable[[Tensor], bool]]): Constraints defining the convex set C.
        """
        self.mirror_map = mirror_map
        if isinstance(step_size, float):
            self.step_size = lambda k: step_size
        elif callable(step_size):
            self.step_size = step_size
        else:
            raise TypeError("step_size must be a float or a callable function of iteration.")
        self.constraints = constraints

    def project(
        self, y: Tensor, constraints: Optional[Callable[[Tensor], bool]] = None
    ) -> Tensor:
        """
        Project tensor y onto the convex set C using the mirror map's projection.

        Args:
            y (Tensor): Tensor to project.
            constraints (Optional[Callable[[Tensor], bool]]): Constraints defining C.

        Returns:
            Tensor: Projected tensor.
        """
        logger.debug("Projecting y onto C using the mirror map's projection.")
        return self.mirror_map.projection(y, constraints)

    def step(
        self,
        x_current: Tensor,
        gradient: Tensor,
        iteration: int
    ) -> Tensor:
        """
        Perform a single Mirror Descent update step.

        Args:
            x_current (Tensor): Current iterate x^(k).
            gradient (Tensor): Gradient or sub-gradient g^(k) of f at x^(k).
            iteration (int): Current iteration number k.

        Returns:
            Tensor: Next iterate x^(k+1).
        """
        eta_k = self.step_size(iteration)
        logger.debug(f"Iteration {iteration}: Step size eta_k = {eta_k}")

        # Compute gradient of Phi at current x
        grad_phi_x = self.mirror_map.grad_phi(x_current)
        logger.debug(f"Gradient of Phi at x_current: {grad_phi_x}")

        # Compute y^{k+1} in the dual space
        grad_phi_y = grad_phi_x - eta_k * gradient
        logger.debug(f"Gradient in dual space (grad_phi_y): {grad_phi_y}")

        # Compute y^{k+1} in the primal space using the conjugate gradient
        y_next = self.mirror_map.grad_phi_star(grad_phi_y)
        logger.debug(f"Primal space y_next before projection: {y_next}")

        # Project y_next onto the convex set C
        x_next = self.project(y_next, self.constraints)
        logger.debug(f"Primal space x_next after projection: {x_next}")

        return x_next

    def optimize(
        self,
        x_init: Tensor,
        objective: Callable[[Tensor], Tensor],
        num_iterations: int,
        verbose: bool = False
    ) -> Tensor:
        """
        Run the Mirror Descent optimization algorithm.

        Args:
            x_init (Tensor): Initial point x^(0) ∈ C.
            objective (Callable[[Tensor], Tensor]): The objective function to minimize.
            num_iterations (int): Number of iterations to perform.
            verbose (bool): If True, print progress.

        Returns:
            Tensor: Optimized point x^(K).
        """
        x = x_init.clone().detach().requires_grad_(False)
        logger.info("Starting Mirror Descent optimization.")
        for k in range(num_iterations):
            # Compute gradient of the objective at current x
            obj = objective(x)
            obj.backward()
            try:
                gradient = x.grad.clone()
            except AttributeError:
                raise RuntimeError("Objective function must be differentiable.")
            x.grad.zero_()
            if verbose and k % max(1, num_iterations // 10) == 0:
                logger.info(f"Iteration {k}: Objective = {obj.item()}")

            # Perform Mirror Descent step
            x = self.step(x, gradient, k)

        logger.info("Mirror Descent optimization completed.")
        return x


def euclidean_projection_simplex(y: Tensor) -> Tensor:
    """
    Project a tensor onto the simplex using the method by Duchi et al. (2008).

    Args:
        y (Tensor): Input tensor.

    Returns:
        Tensor: Projected tensor onto the simplex.
    """
    logger.debug("Projecting onto the simplex.")
    sorted_y, _ = torch.sort(y, descending=True)
    cssv = torch.cumsum(sorted_y, dim=0) - 1
    ind = torch.arange(1, len(y) + 1, device=y.device)
    cond = sorted_y - cssv / ind.float() > 0
    rho = ind[cond][-1]
    theta = cssv[rho - 1] / rho.float()
    return torch.clamp(y - theta, min=0)


class SimplexProjectionConstraint:
    """
    Constraint class for projecting onto the simplex.
    """

    def __call__(self, y: Tensor) -> Tensor:
        """
        Project tensor y onto the simplex.

        Args:
            y (Tensor): Tensor to project.

        Returns:
            Tensor: Projected tensor.
        """
        return euclidean_projection_simplex(y)


def example_squared_l2_norm():
    """
    Example usage of Mirror Descent with Squared L2-norm mirror map.
    Minimizes f(x) = (1/2) * ||x - a||^2, subject to ||x||_2 <= R.
    """
    logger.info("Running example with Squared L2-norm mirror map.")

    # Define parameters
    a = torch.tensor([1.0, 2.0, 3.0])
    R = 5.0

    # Define the objective function f(x) = 0.5 * ||x - a||^2
    def objective(x: Tensor) -> Tensor:
        return 0.5 * torch.sum((x - a) ** 2)

    # Define constraints: ||x||_2 <= R
    def constraint_project(y: Tensor) -> Tensor:
        norm_y = torch.norm(y, p=2)
        if norm_y <= R:
            return y
        else:
            return y * (R / norm_y)

    # Initialize mirror map and optimizer
    mirror_map = SquaredL2NormMirrorMap()
    optimizer = MirrorDescent(
        mirror_map=mirror_map,
        step_size=0.1,
        constraints=constraint_project
    )

    # Initialize x
    x_init = torch.zeros_like(a)

    # Run optimization
    num_iterations = 100
    x_opt = optimizer.optimize(x_init, objective, num_iterations, verbose=True)

    logger.info(f"Optimized x: {x_opt}")
    logger.info(f"Objective value at optimized x: {objective(x_opt).item()}")


def example_negative_entropy():
    """
    Example usage of Mirror Descent with Negative Entropy mirror map.
    Minimizes f(x) = -c^T x, subject to x ∈ simplex.
    """
    logger.info("Running example with Negative Entropy mirror map.")

    # Define parameters
    c = torch.tensor([1.0, 2.0, 3.0])

    # Define the objective function f(x) = -c^T x
    def objective(x: Tensor) -> Tensor:
        return -torch.dot(c, x)

    # Define constraints: x ∈ simplex
    simplex_constraint = SimplexProjectionConstraint()

    # Initialize mirror map and optimizer
    mirror_map = NegativeEntropyMirrorMap()
    optimizer = MirrorDescent(
        mirror_map=mirror_map,
        step_size=0.1,
        constraints=simplex_constraint
    )

    # Initialize x (must be in the simplex)
    x_init = torch.ones_like(c) / len(c)

    # Run optimization
    num_iterations = 100
    x_opt = optimizer.optimize(x_init, objective, num_iterations, verbose=True)

    logger.info(f"Optimized x: {x_opt}")
    logger.info(f"Objective value at optimized x: {objective(x_opt).item()}")


def main():
    """
    Main function to run examples.
    """
    example_squared_l2_norm()
    example_negative_entropy()


if __name__ == "__main__":
    main()