Efficiency-Aware Policy Optimization (EAPO)

• Efficiency-Aware Policy Optimization (EAPO) is a reinforcement learning framework that balances policy performance with computational and sampling costs.
• It employs concave surrogate bounds and iterative optimization to guarantee monotonic or bounded policy improvement while limiting the number of updates.
• Extensions such as multi-agent batching and control variates enhance scalability and efficiency by reducing training and inference time in practical applications.

Efficiency-Aware Policy Optimization (EAPO) refers to a class of methods in reinforcement learning that explicitly address the balance between policy performance and computational (or sampling) cost. These methods seek to maximize the improvement in expected return per unit of computation or data—often under constraints on the number of policy updates, fresh rollouts, or inference time—while guaranteeing monotonic or bounded improvement in performance. EAPO includes both single-agent formulations (with strong ties to off-policy evaluation and surrogate optimization) and multi-agent extensions that harness structural properties of the problem (such as agent independence) to scale efficiently.

1. Foundations and Motivation

In typical policy optimization, achieving high performance may require frequent deployment of new policies and expensive data collection. However, in many practical settings (e.g., robotics, industrial control, large-scale advertising), either collecting new trajectories or redeploying policies is costly or risky. EAPO addresses this by designing optimization algorithms that can obtain significant policy improvements with a limited number of updates or rollouts—maximizing the empirical return while carefully controlling sample and computation usage.

At the core of EAPO is the use of tight surrogate bounds—typically concave (for lower-bounding gains) or convex (for bounding loss in the presence of negative rewards)—that are maximized efficiently given logged data. In the multi-agent context, EAPO also encompasses principled batching structures to minimize the sequential dependencies among agents, thereby improving parallelism and reducing overall training or inference time.

2. Concave Surrogate Bounds and Surrogate Optimization

Let $\theta\in\mathbb{R}^d$ parametrize a stochastic policy $\pi(a|s;\theta)$ in a Markov Decision Process, and consider maximizing the expected total return:

$$J(\theta) = \mathbb{E}_{\tau\sim p(\cdot|\theta)}[R(\tau)] = \int_\tau p(\tau|\theta) R(\tau)\, d\tau$$

Here, $p(\tau|\theta)$ is typically log-concave in $\theta$, but $J(\theta)$ itself is not, due to the trajectory product structure.

Efficiency-aware policy optimization constructs a concave lower bound for $J(\theta)$ using the scalar inequality $x \geq 1 + \log x$. For any anchor parameter $\nu$,

$$p_\nu(\tau|\theta) = p(\tau|\nu)\left[1 + \log\left(\frac{p(\tau|\theta)}{p(\tau|\nu)}\right)\right]$$

is concave in $\theta$ and lower-bounds $p(\tau|\theta)$, with local equality and slope match at $\theta = \nu$.

Given a set of $N$ logged trajectories $\{\tau_i, R(\tau_i)\}_{i=1}^N$ collected under a behavior policy $\theta_0$, the expected return is estimated via importance sampling:

$$\widehat{J}(\theta) = \frac{1}{N}\sum_{i=1}^N R(\tau_i)\frac{p(\tau_i|\theta)}{p(\tau_i|\theta_0)}$$

The surrogate becomes

$$\widehat{J}_\nu(\theta) = \frac{1}{N}\sum_{i=1}^N R(\tau_i) \frac{p(\tau_i|\nu)}{p(\tau_i|\theta_0)}\left[1 + \log\left(\frac{p(\tau_i|\theta)}{p(\tau_i|\nu)}\right)\right]$$

which is concave, tangent to the true objective at $\theta = \nu$, and can be efficiently optimized as a proxy for $J$.

When rewards $R(\tau_i)$ are negative or mixed in sign, EAPO incorporates a convex upper bound for those samples, constructing a piecewise surrogate that remains concave in $\theta$.

3. Iterative Optimization Algorithms

The central EAPO algorithm for single-agent cases—sometimes referenced as iPoWER—proceeds by iteratively maximizing surrogates built from the most recent policy parameters $\theta$. The update loop typically consists of:

1. Building a locally tight concave surrogate $\widehat{J}_\theta(\varphi)$ at the current parameter $\theta$.
2. Maximizing this surrogate with respect to $\varphi$ using a concave optimization solver (such as L-BFGS or Hessian-free Newton).
3. Setting $\theta\leftarrow \varphi^*$ and repeating for $T$ inner iterations before any new data collection.

Pseudocode representation:

for t in range(T): # Build concave surrogate at current θ def surrogate(ϕ): return (1/N) * sum_over_i( R(τ_i) * (p(τ_i|θ)/p(τ_i|θ₀)) * (1 + log(p(τ_i|ϕ)/p(τ_i|θ))) ) # Concave maximization step θ = argmax_ϕ surrogate(ϕ)

Convergence to a stationary point is guaranteed under standard compactness and smoothness assumptions, and each optimization pass monotonically improves the empirical objective due to the surrogate's local tightness properties.

4. Computational Efficiency and Sample Utilization

The per-iteration computational cost for the surrogate maximization is dominated by evaluating the surrogate and its gradient on $N$ samples, with complexity $O(Nd)$ per gradient evaluation and $O(kNd)$ per inner optimization step (where $k$ is the number of solver steps). Because new rollouts are far more expensive than on-dataset optimization, re-using the same fixed dataset for $T$ inner updates provides a significant reduction in policy deployment cost.

Empirical results indicate that a single batch of $N\approx 10^8$ to $10^9$ samples and $T\approx 20$ inner updates matches or exceeds the performance of $T$ full redeployments, drastically raising sample efficiency. Importance-weight clipping and strong control variates (baselines) are critical for variance reduction, especially as $T$ increases.

5. Extensions: Negative Rewards and Control Variates

The surrogate construction for EAPO inherently assumes nonnegative rewards. To accommodate trajectories with $R(\tau_i)<0$, the algorithm uses a convex upper-bound for the importance term:

$$u_\nu(\tau|\theta) = p(\tau|\nu)\exp\left[(\theta-\nu)^\top \nabla_\theta \log p(\tau|\theta)|_{\theta=\nu}\right]$$

The piecewise surrogate employs $1 + \log(\cdot)$ for $R\geq 0$ and $\exp(\cdot)$ for $R < 0$, maintaining global concavity.

A baseline $b$ can be introduced as a control variate, subtracting $b$ from each reward and adding $b\, E[p(\tau|\theta)/p(\tau|\theta_0)]$, to reduce gradient variance. Optimal baseline coefficients are proportional to covariance and variance terms; in practice, a fraction (0.5–0.99) of the optimal value often best controls bias-variance trade-off.

6. Multi-Agent Efficiency-Aware Policy Optimization

In multi-agent reinforcement learning, EAPO principles are instantiated with mechanisms that allow partial parallelism, balancing efficiency and the compositional monotonicity of policy improvement. The B2MAPO (Batch-by-Batch Multi-Agent Policy Optimization) algorithm explicitly constructs disjoint batches of agents—grouped by learned dependency graphs—where each batch is updated in sequence but agents within a batch are updated simultaneously.

The batch structure is generated in a two-layer hierarchy:

• Upper Layer: Uses encoded trajectory summaries to construct an agent-dependency attention graph. Weak dependencies are thresholded, forming a DAG whose topological sort defines update batches.
• Lower Layer: For each batch, a clipped surrogate objective (with off-policy and batch correction) is maximized for that batch’s parameters. In parallel, a distilled individual policy is trained with cross-regularization, providing efficient execution-time policies.

Performance loss due to batching is rigorously upper-bounded via local total-variation and advantage discrepancies. By controlling batch size $|B|$ and DAG threshold $\delta$, users directly trade off efficiency (smaller $|B|$ = more parallelism) against tightness of the improvement bound.

7. Empirical Results and Trade-offs

EAPO methods, both in single and multi-agent contexts, have demonstrated substantial empirical gains in sample and wall-clock efficiency. Single-agent iPoWER experiments on Gym Cartpole show that increasing inner iteration count $T$ from 1 to 5 yields markedly higher sample efficiency, with diminishing returns beyond $T\approx 5$–20 due to variance. Strong control variates are essential for exploiting the full surrogate improvement. In large-scale online advertising, iPoWER with $T\approx 50$ provides up to 60-fold improvement over a single policy update.

B2MAPO achieves up to 60.4% reduction in training time and 78.7% reduction in execution time relative to fully sequential agent-wise optimization (A2PO) in StarCraftII challenges, while matching or exceeding win rates. Hyperparameter tuning of batch size, update period, and surrogate clipping is essential for optimizing the efficiency–performance trade-off.

8. Theoretical Guarantees and Generalization

EAPO establishes provable monotonic improvements or explicit bounds on suboptimality for each policy update sequence. In the single-agent case, monotonic ascent of the empirical return is guaranteed. In the multi-agent setting, a tight telescopic bound over batchwise policy total-variation quantifies the degradation, providing a continuum between fully parallel (efficient but looser) and fully sequential (tighter but less efficient) updates.

Generalization of EAPO entails the following canonical steps:

1. Learning or modeling a dependency graph among subcomponents (agents, modules).
2. Optimizing in parallelizable batches, subject to performance degradation bounds.
3. Maintaining a distillation pathway from a “full” policy to a lightweight deployable surrogate.
4. Providing (and monitoring) theoretical performance-cost bounds of the form $$|J(\pi_{\text{new}}) - \text{Surrogate}(\pi_{\text{old}}, \pi_{\text{new}})| \leq \text{CostBudget}\cdot\epsilon_{\text{bound}}$$.

These guarantees, along with empirical success, support EAPO as a practical and theoretically principled approach for resource-constrained policy learning.