Density-Matrix Reinforcement Learning

Chiuman Ho

cs.LGOct 12, 2026

Density-Matrix Reinforcement Learning

Authors: Chiuman Ho

Organizations: TetraMem Inc.

Abstract

What information does a reinforcement learning policy omit? At each state, a standard policy stores one probability per action. It has no place to record how actions relate to each other, such as which pairs go together and which conflict. We ask when this missing structure matters, and we give exact answers. We store the policy at each state as a density matrix, a positive semidefinite, unit-trace matrix borrowed from quantum information theory. We use it here as a purely classical data structure, computed on ordinary classical hardware. The diagonal holds the usual action probabilities. The off-diagonal entries, called coherences, are new parameters. This gives density-matrix reinforcement learning (DMRL). Coherences change outcomes in two settings. First, suppose a policy is regularized toward a coherent reference policy using matrix relative entropy. The optimum then has coherences that no diagonal reference policy produces, and its action probabilities can depend on the reference policy's coherences. This is relevant to fine-tuning of large language models (LLMs) with a Kullback--Leibler penalty. Second, suppose the agent chooses a unit-vector action and the reward is a quadratic function of that vector. This could arise when an LLM judge scores a weighted consolidation of reasoning paths, or when coupled robotic sensors each respond to a mix of actions. Then some coherent policies earn strictly more expected reward than every classical policy with the same action probabilities, and we give the margin in closed form. In a two-action example, both policies assign probability $0.5$ to each action, yet the classical one earns expected reward $0$ while the coherent one earns $1$ , the maximum. Outside these two settings, DMRL matches or transfers the classical results. This is a theory paper. All results are proved, and we report no empirical evaluation.

Figures & tables

Operator concept	Classical analogue (diagonal case)
Density matrix $\rho$	Probability vector
Diagonal $(\rho_{s})_{ii}$	Action probabilities $\pi(a_{i}\mid s)$
Coherence $\rho_{ij}$ , $i\neq j$	None (zero in the classical case)
Observable $O$ ; expectation $\Tr(\rho O)$	Reward table; expected value $\sum_{i}\pi_{i}O_{ii}$
Von Neumann entropy $S(\rho)$	Shannon entropy $-\sum_{i}\pi_{i}\log\pi_{i}$
Matrix relative entropy $D(\rho\\|\rho^{\prime})$	KL divergence $D_{\mathrm{KL}}(\pi\\|\pi^{\prime})$

Table 1: Operator concepts and their classical analogues. Setting the off-diagonal entries of

\rho

(and of

\rho^{\prime}

and

O

, where relevant) to zero recovers the right-hand column. The diagonal is unchanged by this operation. The coherences have no classical counterpart because they vanish.

	Ordinary RL	Blended-action RL
Action	one label $a_{i}$	one unit vector $\mathbf{x}\in\mathbb{R}^{k}$
Reward	$R_{i}$ , one number per label	$R(\mathbf{x})=\mathbf{x}^{\top}O\mathbf{x}$
Policy statistic	$p_{i}=\pi(a_{i})$	$\rho=\mathbb{E}_{\pi}[\mathbf{x}\mathbf{x}^{\top}]$

Table 2: Two action models. The standard basis vectors

\mathbf{e}_{i}

are unit vectors, and choosing

\mathbf{x}=\pm\mathbf{e}_{i}

earns

\mathbf{x}^{\top}O\mathbf{x}=O_{ii}

, the reward of action

a_{i}

. So an agent that chooses only

\pm\mathbf{e}_{1},\dots,\pm\mathbf{e}_{k}

is exactly the left column. This restriction is the classical baseline used throughout this subsection.

Reinforcement Learning

Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.

See all 1,727 papers

Sep 30, 2026

Make Sparse Rewards Count: Density-Aware Reward Aggregation for Multi-Reward RL

Multi-reward reinforcement learning trains large language models to satisfy multiple behavioral objectives simultaneously. Reward-wise normalization, as used in GDPO, preserves reward-specific relative information within rollout groups, but different objectives can still exhibit uneven learning progress. We study this behavior through advantage energy, the sum of a reward's squared advantages over a batch. Under idealized GDPO normalization, we show that this energy is proportional to active-group density: the fraction of rollout groups in which the reward provides nonzero relative advantages. This reveals a residual batch-level signal imbalance and provides a basis for calibrating reward contributions. Based on this relation, we propose Density-Aware Reward Aggregation (DARA). We derive an inverse-square-root density correction that gives greater weight to signals from less frequently active rewards. DARA computes its weights from each rollout batch, adapting to changes in reward activity throughout training without modifying the underlying policy optimization objective. Experiments on tool calling and mathematical reasoning show that DARA learns the targeted behaviors faster than GDPO, reaching high format compliance in up to 26% fewer training steps on tool calling and near-saturated length compliance in up to 65% fewer steps on mathematical reasoning, while remaining competitive in final performance. Our code is available at https://github.com/zhaihaotian/DARA.

Tong Zheng, Skylar Zhai, Zhan Cheng +7

University of Chinese Academy of Sciences · University of Minnesota Twin Cities · University of Wisconsin–Madison +5

Multi-Objective Reinforcement Learning Reinforcement Learning RL for Language Model Reasoning Sparse-Reward RL Policy Learning

Ask PDF

Sep 27, 2026

Diffusion Reward Models

Reward models underpin the alignment of large language models, yet the dominant designs reduce each prompt--response pair to a point estimate or to a distribution from a fixed parametric family. This is at odds with human preference, which is inherently multimodal: the same response can be reasonably judged in many ways, and no single family covers all of them. To better fit this structure, we introduce DRM, a Diffusion Reward Model that recasts reward modeling as conditional density estimation over

p(\mathbf{r}\mid x,y)

. Conditioned on a frozen LLM encoder, a lightweight Diffusion Transformer denoises Gaussian noise into a reward vector, placing no parametric assumption on the output distribution and naturally representing its multimodal structure. A single architecture handles both multi-attribute regression and pairwise preference data, and at inference

N

samples form an empirical reward distribution that can be aggregated into a scalar, a variance, or quantiles. Across five benchmarks, DRM matches or surpasses baselines under matched data and backbone, stays competitive with much larger discriminative, distributional, and generative RMs despite its modest training scale, and recovers multimodal reward structure where conventional heads collapse to a point. Uncertainty-aware rejection and lower-confidence-bound (LCB) aggregation further demonstrate that DRM can exploit distributional information beyond a scalar reward to improve reward-model decisions. Downstream RLHF experiments additionally show that using DRM as the training-time reward leads to improved policy performance, directly validating the practical benefit of diffusion-based reward modeling for RLHF training.

Xiangyang Wang, Bingxiang He, Zeyuan Liu +13

Tsinghua University · University of Illinois Urbana-Champaign

Reward Modeling Multimodal Reward Modeling Diffusion Models RL from Human Feedback

Ask PDF

Jun 6, 2026

QnRL: Quantum-Native Reinforcement Learning

Quantum reinforcement learning (QRL) is a promising approach to learn effective decision strategies across several applications with stochastic environments. Instead of directly modeling the random variables that govern these environments, existing QRL architectures indirectly approximate environment behavior by estimating expected outcomes, which limits their expressive power and adaptive potential. Overcoming such challenges requires a novel QRL approach that exploits the distributional nature of quantum computers to directly model environment random variables as quantum state distributions. Hence, in this paper, a novel framework dubbed quantum-native reinforcement learning (QnRL) is proposed. QnRL is a distributional RL framework that learns conditional distributions naturally in Hilbert space via superimposed and entangled quantum states. Thus, QnRL can directly model the behavior of stochastic learning environments via the natural properties of quantum systems. QnRL accomplishes this via a novel, proposed quantum amplitude kickback (QuAK) algorithm that enables comparing the

n

-th power of the

m

-th moment of multiple superimposed distributions. It is theoretically proven that a conditional action policy distribution is distilled from the moments of a quantum generative model entirely within Hilbert space via QuAK, and optimized via QnRL. This complex distribution composition is also shown to provide extra dimensions for expressing environment correlations that are unknown to purely classical and classically-sampled quantum distributional models. Experimental results across diverse environments show that QnRL achieves up to

82.9\%

higher evaluation scores, with up to

94.3\%

fewer parameters on average, more accurately estimates the expected return for unseen observations, and better adapts to varying stochastic conditions compared to the baseline.

Alexander DeRieux, Walid Saad

Bradley Department of Electrical and Computer Engineering, Virginia Tech Institute for Advanced Computing, Alexandria, VA 22305 USA

Distributional RL Quantum Generative Modeling Quantum Machine Learning Quantum RL

Ask PDF

Abstract

Reinforcement Learning

Explore similar work

Make Sparse Rewards Count: Density-Aware Reward Aggregation for Multi-Reward RL

Diffusion Reward Models

QnRL: Quantum-Native Reinforcement Learning