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 : develop code below Mirror Descent in great details using torch module only from stratch 


2.2 Mirror Descent
Mirror descent (MD) is a general framework for the construction and analysis of optimization
algorithms [Nemirovski and Yudin, 1983]. Its key machinery is a pair of conjugate mirror maps that
map the iterates of an optimization algorithm back-and-forth between a primal space and a dual space.
We follow the common practice of defining the mirror maps with the gradient mapping of a convex
function of Legendre-type [Rockafellar, 1970, Section 26].
Let’s first define Bregman divergence and Bregman projection. Suppose that Ψ is a convex function
of Legendre type. It induces a Bregman divergence between any x ∈ dom Φ and y ∈ int (dom Φ):
DΦ(x, y) = Φ(x) − Φ(y) − ⟨∇Φ(y), x − y⟩. (5)
Let C ∈ dom Φ be a closed convex set. The Bregman projection of any y ∈ int (dom Φ) onto C is
projΦ
C
(y) = arg min
x∈C
DΦ(x, y). (6)
Properties of Bregman projection can be found in, e.g., Bauschke and Borwein [1997].
2
Now consider the problem of minimizing a convex function f : X → R ∪ {∞} over a closed convex
set C ⊂ dom Φ. We use the presentation of MD given by Bubeck [2015]: at each iteration k,
1. Given x
(k)
, find y
(k+1) such that
∇Φ(y
(k+1)) = ∇Φ(x
(k)
) − ηkg
(k)
. (7)
where ηk is the step size and g
(k)
is the gradient ∇f(x
(k)
) or a sub-gradient of f at x
(k)
.
2. Compute x
(k+1) = projΦ
C
(y
(k+1)).
Define the conjugate function of Φ as Φ
∗
(x
∗
) = supx∈dom Φ {⟨x, x∗
⟩ − Φ(x)}. Then, using the
definition in (6) and the identity ∇Φ
∗
(∇Φ(x)) = x, we can express it more compactly as
x
(k+1)= arg min
x∈C
DΦ

x, ∇Φ
∗

∇Φ(x
(k)
)−ηkg
(k)


, (8)
which can be further simplified to [Beck and Teboulle, 2003, Bubeck et al., 2012]
x
(k+1) = arg min
x∈C
n
ηk


g
(k)
, x

+ DΦ(x, x(k)
)
o
. (9)
Next we discuss three examples of the MD algorithm for solving minx∈∆ f(x), where ∆ is the
simplex. Each leads to a variant of the DAPO method we will present in Section 3.
Example 2.1 (Squared L2-norm). Let Φ(x) = 1
2
∥x∥
2
2
, which is Legendre type with int (dom Φ) =
dom Φ = R
n. We have Φ
∗
(x
∗
) = 1
2
∥x
∗∥
2
2
, ∇Φ(x) = x, ∇Φ
∗
(x
∗
) = x
∗
, and DΦ(x, y) =
1
2
∥x − y∥
2
2
. In this case, the MD algorithm (8) becomes the classical projected gradient method
x
(k+1) = arg min
x∈∆


x −

x
(k) − ηk∇f(x
(k)
)


2
2
.
Example 2.2 (Negative entropy on R
n
+). Consider the negative entropy Φ(x) = P
i
(xi
log(xi) − xi)
with dom Φ = R
n
+ (and the convention 0 log 0 = 0). It is of Legendre type, with Φ
∗
(x
∗ P
) =
i
exp(x
∗
i
), ∇Φ(x) = log(x) and ∇Φ
∗
(x
∗
) = exp(x
∗
), where log and exp apply component-wise
to vectors. For any x ∈ R
n
+ and y ∈ R
n
++, their Bregman divergence is the KL-divergence:
DΦ(x, y) = P
i
(xi
log(xi/yi) − xi + yi). (10)
In this case, the Bregman projection of y ∈ R
n
++ onto ∆ is projΦ
∆(y) = y/∥y∥1 and (8) becomes
x
(k+1) = x
(k)
exp(−ηkg
(k)
)/∥x
(k)
exp(−ηkg
(k)
)∥1. (11)
Example 2.3 (Negative entropy on ∆). Let ϕ(x) = P
i
(xi
log(xi) − xi) and define Φ(x) =
ϕ(x) + δ(x|∆), where δ(·|∆) is the indicator function of ∆, i.e., δ(x|∆) = 0 if x ∈ ∆ and +∞
otherwise. Apparently dom Φ = ∆, which has an empty interior. As a result, Φ is not of Legendre
type and in fact is not differentiable (see Appendix B). However, the MD algorithm is still welldefined. Specifically, in (7) we interpret ∇Φ(x
(k)
) as any subgradient in the subdifferential ∂Φ(x
(k)
),
and find y
(k+1) such that there exists some ∇Φ(y
(k+1)) ∈ ∂Φ(y
(k+1)) to make the equality hold.
Despite ∇Φ(y) being multi-valued as a subgradient, the Bregman divergence (5) is still well defined
(Corollary B.3). As a result, DΦ is the same as (10). Using the fact x, y ∈ ∆, it can be simplified as
DΦ(x, y) = P
i
xi
log(xi/yi). (12)
In addition, we have Φ
∗
(x
∗
) = log (P
i
exp(x
∗
i
)) with dom Φ∗ = R
n [Rockafellar, 1970, Section 16]. Clearly, Φ
∗
is a differentiable function throughout R
n and
∇Φ
∗
(x
∗
) = exp(x
∗
)/∥ exp(x
∗
)∥1.
In this case, the MD algorithm (8) yields the same update as (11). However, the projection step is no
longer needed because the range of ∇Φ
∗
is the interior of ∆ and we can simply express MD as
x
(k+1) = ∇Φ
∗

∇Φ(x
(k)
) − ηkg
(k)

.
Remark 2.4. Although Examples 2.2 and 2.3 give the same update (11), there are subtle differences
in the theory. In particular, we have ∇Φ
∗ = (∇Φ)−1
in Example 2.2 and the equivalence between (8)
and 9 is easily established. However, this is not the case for Example 2.3 because Φ is not Legendre
and nondifferentiable. Consequently, we can no longer leverage the convex optimization machinery
as in the tabular case [e.g., Xiao, 2022, Lan, 2023] for convergence analysis with general function
approximation. Instead, we extend the classical MD theory to work with mirror maps whose gradient
mapping and conjugate mapping are not inverses of each other; see Appendix B, Lemma B.2.
Algorithm 1 Dual Approximation Policy Optimization (DAPO)
1: Input: Initialize policy π
(0) with parameters θ
(0); mirror map Φ
2: for k = 0, . . . , K − 1 do
3: Find Qb(k)
that approximates Q(k)
(Critic Update)
4: Find θ
(k+1) that (approximately) solves the problem
min
θ∈Θ
Es∼d
(k)
ρ
h
DΦ∗

∇Φ(π
(k)
s
) − ηkQb(k)
s
, f θ
s

i
5: Assign π
(k+1)
s = projΦ
∆(A)

∇Φ
∗
(f
θ
(k+1)
s
)

, s ∈ S
6: end fo