Improving Reward-Conditioned Policies for Multi-Armed Bandits using Normalized Weight Functions

TL;DR

This paper identifies a key limitation of reward-conditioned policies (RCPs) in multi-armed bandits—slow convergence and suboptimal rewards in discrete-action settings. It introduces generalized marginalization, where inference policies are formed via a normalized weight function $w_\phi(r)$ (with $\int w_\phi(r) dr = 1$ or $\sum_r w_\phi(r)=1$) that may take negative values, enabling more reward-focused and distinct action distributions $\pi_{\theta,\phi}^{\ddagger}(a)$. Two practical strategies are proposed—Optimistic and SubMax—and an optimization-based variant that maximizes the expected reward; these are evaluated across non-contextual, contextual, and combinatorial MABs using reward-conditioned generative models (CVAEs). Empirical results show that SubMax frequently achieves superior performance and robustness, particularly in large action spaces and sparse reward regimes, making RCPs competitive with classical methods such as UCB and Thompson sampling. The work suggests promising directions for extending RCPs to real-valued actions and continual learning, potentially broadening their applicability to real-world design and optimization tasks.

Abstract

Recently proposed reward-conditioned policies (RCPs) offer an appealing alternative in reinforcement learning. Compared with policy gradient methods, policy learning in RCPs is simpler since it is based on supervised learning, and unlike value-based methods, it does not require optimization in the action space to take actions. However, for multi-armed bandit (MAB) problems, we find that RCPs are slower to converge and have inferior expected rewards at convergence, compared with classic methods such as the upper confidence bound and Thompson sampling. In this work, we show that the performance of RCPs can be enhanced by constructing policies through the marginalization of rewards using normalized weight functions, whose sum or integral equal $1$, although the function values may be negative. We refer to this technique as generalized marginalization, whose advantage is that negative weights for policies conditioned on low rewards can make the resulting policies more distinct from them. Strategies to perform generalized marginalization in MAB with discrete action spaces are studied. Through simulations, we demonstrate that the proposed technique improves RCPs and makes them competitive with classic methods, showing superior performance on challenging MABs with large action spaces and sparse reward signals.

Figures (8)

• Figure 1: Illustration of different strategies to construct inference policies. Figure \ref{['fig:illustration-env']} shows the true value per action and the rest of the figures show different inference-time policies where the red line indicates the ground truth best arm. Numbers in braces are the expected reward (larger the better) under each policy and the number next to the environment corresponds to a random policy. Note that the 1-conditioned policy, or positive policy, is equivalent to the optimistic policy.
• Figure 2: Accumulated regret with varying $K$ and $(\alpha, \beta) = (1, 9)$ and $N_b = 500$.
• Figure 3: Accumulated regret with varying $(\alpha, \beta)$ while keeping $K = 100$ and $N_b = 100$.
• Figure 4: Accumulated regret with varying $N_b$ while keeping $K = 100$ and $(\alpha, \beta) = (1, 9)$.
• Figure 5: Accumulated regret with varying constant shift $s$ in the ground truth value function. The numbers next to each $s$ in the subtitles are the average of the expected reward over all randomized environment.