Adaptive Action Weighting (AAW)

• Adaptive Action Weighting (AAW) is a framework that dynamically learns weighting functions for actions, events, or tasks to determine their credit in sequential and meta-learning contexts.
• It replaces fixed heuristics like the λ-return with parameterized, context-dependent weight functions optimized via meta-gradient methods and trajectory optimization.
• Empirical evaluations show AAW accelerates convergence and improves performance in RL and few-shot classification through enhanced temporal credit assignment and adaptive task weighting.

Adaptive Action Weighting (AAW) is a principled framework for dynamically learning weighting schemes over actions, events, or tasks in sequential or meta-learning problems. The approach replaces fixed or hand-designed assignment coefficients with parameterized, learnable weight functions, optimizing these weights online via meta-gradient methods or trajectory optimization. AAW generalizes conventional credit assignment heuristics, offering improved temporal credit assignment, task weighting, and accelerated learning dynamics in reinforcement learning (RL) and meta-learning contexts (Zheng et al., 2021, Nguyen et al., 2023).

1. Conceptual Foundations

AAW formalizes the assignment of credit (or blame) by introducing a parameterized weighting function that adaptively determines the influence of one entity (such as an action or task) on an outcome (e.g., reward, loss) at a later time. In RL, this replaces fixed heuristics such as the λ-return, where credit decays exponentially by a scalar λ, with a learnable function $w_\eta$ that can depend on the full context—states, actions, outcomes, and time lags. In meta-learning, AAW treats the vector of task weights within each meta-update as an action in a trajectory optimization setting, where the system state is the meta-learner’s parameter vector (Zheng et al., 2021, Nguyen et al., 2023).

2. Methodological Formulation

In RL, AAW defines a scalar pairwise-weight function $w_{\eta}(s_i, s_j, j-i) \in [0,1]$ that quantifies the contribution of the action at $s_i$ (time $i$) to the reward or TD-error at $s_j$ (time $j$). The function is parameterized by $\eta$ and can be instantiated either through embedding states and lags followed by multiplicative fusion and a multilayer perceptron with a sigmoid output, or as a tabular parameterization in small domains.

Key return and TD estimators include:

• Pairwise-Weighted Return (PWR):

$$\hat\Psi^{\mathrm{PWR}}_t = \left( \sum_{j=t+1}^T w_\eta(s_t, s_j, j-t) R_j \right) - v(s_t)$$

• Pairwise-Weighted TD-error (PWTD):

$$\hat\Psi^{\mathrm{PWTD}}_t = \sum_{j=t+1}^T w_\eta(s_t, s_j, j-t) \delta_j$$

Policy gradient updates are then performed with these estimators, generalizing the standard $\lambda$-return, which is recovered as a special case with $w_{\eta}(s_i, s_j, j-i) \in [0,1]$0 (Zheng et al., 2021).

In meta-learning, AAW (under the TOW framework) formulates task weighting as a control action. At each meta-update step $w_{\eta}(s_i, s_j, j-i) \in [0,1]$1, the learner selects a weight vector $w_{\eta}(s_i, s_j, j-i) \in [0,1]$2 over $w_{\eta}(s_i, s_j, j-i) \in [0,1]$3 tasks to minimize the running cost

$w_{\eta}(s_i, s_j, j-i) \in [0,1]$4

trading off rapid loss reduction against deviation from uniform weighting. The update dynamics

$w_{\eta}(s_i, s_j, j-i) \in [0,1]$5

define the state evolution, and trajectory optimization over a finite horizon is solved by iterative Linear Quadratic Regulator (iLQR) with local linearization and quadraticization. Convergence to an $w_{\eta}(s_i, s_j, j-i) \in [0,1]$6–stationary point is established under mild conditions on losses and gradients (Nguyen et al., 2023).

3. Meta-Gradient Optimization

AAW deploys a two-loop meta-gradient optimization process for its weight parameters. The inner loop updates the base learner’s parameters (e.g., policy or meta-parameter $w_{\eta}(s_i, s_j, j-i) \in [0,1]$7) using the current weight function and corresponding (PWTD or PWR) advantage estimates. The outer loop evaluates performance under the updated parameter and computes gradients with respect to the weight parameters via the chain rule:

$w_{\eta}(s_i, s_j, j-i) \in [0,1]$8

This yields practical meta-gradient estimates for updating $w_{\eta}(s_i, s_j, j-i) \in [0,1]$9 via gradient ascent (Zheng et al., 2021).

In the TOW meta-learning setting, iLQR backward and forward passes compute optimal per-timestep action corrections and value function derivatives, supporting efficient optimization over weight sequences (Nguyen et al., 2023).

4. Algorithmic Implementation

The typical AAW algorithm alternates between sampling trajectories, computing pairwise-weighted advantage estimates using the current weights, performing inner-loop parameter updates, evaluating these in outer-loop rollouts using standard advantage estimates, and finally updating the weight function's parameters with the meta-gradient. Pseudocode details are provided explicitly in (Zheng et al., 2021). In meta-learning, the TOW algorithm implements a line-searched iLQR trajectory optimization at each meta-iteration.

Domain	Weighting Mechanism	Optimization Algorithm
RL	Pairwise state/time weighting	Meta-gradient ascent
Meta-learning	Task weights per minibatch	iLQR trajectory optim.

5. Empirical Evaluation

AAW demonstrates consistent improvements across RL credit assignment, meta-learning, and few-shot classification:

• RL Tabular Depth-DAGs: Learned weights closely approximate true credit assignments and accelerate convergence by factors of 2–10× compared to fixed $s_i$0 (Zheng et al., 2021).
• Key-to-Door Gridworld: Meta-PWR matches handcrafted credit and outperforms state-of-the-art automated baselines. Meta-PWTD also outperforms best fixed $s_i$1 and other modern methods.
• bsuite Credit Assignment Domains: Meta-PWR/PWTD achieve lowest regret in most domains.
• Atari 49 Games: Meta-PWTD outperforms A2C ($s_i$2) in 30/49 games, ties in 5, with no major regressions.
• Few-shot Classification (TOW/AAW in MAML, ProtoNet): On Omniglot 5way 1shot, MAML+AAW and ProtoNet+AAW exceed uniform, exploration, and exploitation baselines by 0.6–1.6% in accuracy; on mini-ImageNet gains are ~2–3%. AAW adds computational overhead (about 7× slower than uniform weighting), but yields robust benefits (Nguyen et al., 2023).

6. Theoretical Properties and Guarantees

AAW’s design admits theoretical convergence guarantees. In meta-learning, TOW’s iLQR-based updates ensure that, under boundedness and Lipschitz continuity of losses and their derivatives, the method achieves expected squared gradient norm bounded by $s_i$3 plus a decaying term in $s_i$4, the number of meta-updates. As the iLQR action is constrained near the uniform prior and the number of tasks grows, $s_i$5, guaranteeing near-stationarity (Nguyen et al., 2023).

7. Context and Significance

AAW unifies and generalizes fixed-weight heuristics in sequential credit and task weighting problems, subsuming the $s_i$6-return and other exponential discount schemes as special parameter choices. Its deployment in RL and meta-learning shows empirical advantages in sample efficiency, regret minimization, and final accuracy, demonstrating the practical utility of learnable, context-dependent weighting in long-horizon and task-imbalanced optimization. The approach amplifies the flexibility and adaptivity of policy optimization, meta-learning, and temporal credit assignment (Zheng et al., 2021, Nguyen et al., 2023).