Reward-Balanced Batching: Principles and Applications

• Reward-balanced batching is a design principle that dynamically redistributes reward signals within batches to balance heterogeneous criteria.
• In rubric-based reinforcement learning, it recalibrates scalar rewards via focal weights computed from headroom measures, improving hard-criterion performance.
• The approach extends to batched bandits and service systems, optimizing trade-offs such as latency versus energy and reducing overall regret.

Reward-balanced batching is a family of batched optimization and decision procedures in which the effective signal carried by a batch is deliberately reallocated rather than treated as a fixed, uniform aggregate. In the literature, the term appears most explicitly in rubric-based reinforcement learning for LLMs, where it denotes ensuring that, within and across training batches, the policy does not over-optimize “easy” reward dimensions while neglecting harder or more important ones (Huang et al., 26 May 2026). Closely related formulations appear in reinforcement learning from bagged reward, where a bag-level return is redistributed across the transitions inside the bag (Tang et al., 2024); in batched bandits, where batch schedules are chosen to balance regret against limited feedback (Kalkanli et al., 2021); and in service systems, where dynamic batching balances latency and energy under size-dependent batch processing (Xu et al., 4 Jan 2025).

1. Conceptual scope and recurring structure

Across these settings, reward-balanced batching is not a single algorithm but a recurring design principle: a batch is treated as a constrained container of optimization signal, and the learner or controller must decide how that signal should be distributed across criteria, transitions, arms, or jobs. The common concern is that naive batching can make one component dominate the update: an already-saturated rubric criterion in LLM alignment, a late transition in sparse-reward RL, a currently over-sampled arm in batched exploration, or a throughput-favoring batch size in online services.

A concise way to compare the main uses in the cited literature is to distinguish what is being balanced and what object is adapted.

Setting	Balanced quantity	Representative mechanism
Rubric-based RL	Criterion-level reward signal	Saturation-aware focal reweighting
RL from bagged reward	Conserved bag reward across steps	Reward redistribution with bag-sum supervision
Batched bandits	Regret or information per batch	Adaptive or analytically designed batch schedules
Batch services	Latency vs. power/efficiency	SMDP-based dynamic batching

A plausible implication is that the phrase names a structural property rather than a single formalism. What remains constant is the rejection of static aggregation when the components being aggregated have heterogeneous headroom, heterogeneous uncertainty, or heterogeneous operational cost.

2. Rubric-based reinforcement learning and focal reweighting

In rubric-based RL for LLMs, responses are scored along $K$ criteria $c^{(1)},\dots,c^{(K)}$. For a pair of rollouts $y_i$ and $y_j$, the judge returns a criterion-level score vector

$$\mathbf{s}_{i,j} = J(x,y_i,y_j;\mathcal{C}) \in [0,S_{\max}]^K,$$

with scalar aggregation under a static base weight vector $\mathbf{w}_{\text{base}}$ given by

$$S_{i,j}(\mathbf{w}) = \langle \mathbf{w}, \mathbf{s}_{i,j}\rangle = \sum_{k=1}^K w^{(k)} s_{i,j}^{(k)}.$$

The pairwise margin is

$$\Delta_{i,j}(\mathbf{w}) = S_{i,j}(\mathbf{w}) - S_{j,i}(\mathbf{w}),$$

which is passed through a thresholded sign function $\Phi_\tau$, and, in GRPO/GSPO-style group-relative training with $G$ rollouts, rollout $c^{(1)},\dots,c^{(K)}$0 receives scalar reward

$c^{(1)},\dots,c^{(K)}$1

The motivating failure mode is reward polarization: different criteria saturate at different rates, so fixed $c^{(1)},\dots,c^{(K)}$2 lets easy dimensions such as fluency, politeness, verbosity, or format dominate the scalar margin even when harder or task-critical dimensions remain weak. Reward-balanced batching, in this formulation, means recalibrating the scalarizer inside each rollout group so that under-satisfied criteria receive more optimization mass.

Focal Reward implements that recalibration in three stages. First, it computes base rewards $c^{(1)},\dots,c^{(K)}$3 and converts them into Gibbs weights

$c^{(1)},\dots,c^{(K)}$4

Second, for each criterion $c^{(1)},\dots,c^{(K)}$5, it computes rollout-level average scores

$c^{(1)},\dots,c^{(K)}$6

and defines the saturation degree

$c^{(1)},\dots,c^{(K)}$7

Here $c^{(1)},\dots,c^{(K)}$8 is a reward-weighted estimate of frontier saturation: if high-reward rollouts are already strong on criterion $c^{(1)},\dots,c^{(K)}$9, then $y_i$0 is high; if even the frontier is weak on $y_i$1, then $y_i$2 remains low.

Third, it converts headroom $y_i$3 into focal weights: $y_i$4 The final scalar reward is then recomputed with $y_i$5: $y_i$6

The theoretical analysis casts scalarization as a local credit-assignment problem. For a criterion-wise score-difference vector $y_i$7 and latent beneficial update direction $y_i$8, the misallocation probability is

$y_i$9

and under a mild sub-Gaussian residual assumption it is upper-bounded by

$y_j$0

Under symmetric residuals and isotropic simplification, the preferred direction is proportional to a criterion-level signal $y_j$1, and under a headroom-shaped edge model this leads to

$y_j$2

Focal Reward operationalizes this by using $y_j$3 as an observable proxy for frontier headroom and $y_j$4 as a focal-strength parameter (Huang et al., 26 May 2026).

3. Batchwise realization and empirical pattern in rubric-based RL

Focal Reward is already a per-group batched procedure. For each prompt’s rollout group inside a batch, the algorithm computes criterion scores $y_j$5, base rewards, Gibbs weights, saturation degrees, and a group-specific focal scalarizer before computing group-relative advantages. The effective criterion weights are therefore not global constants: they are recomputed online for every group of rollouts, making the reward aggregation simultaneously training-state dependent and prompt dependent.

The same construction also induces a natural batch-level variant. If one wants smoothing across prompts, the literature describes aggregating saturation over the whole batch: $y_j$6 or, equivalently, by a global softmax over all rollouts in the batch. Batch-level focal weights are then defined by

$y_j$7

followed by normalization. This preserves the central idea: the batch itself decides which rubric dimensions still have room for improvement.

The empirical picture is correspondingly specific. Focal Reward was evaluated on three backbones—Qwen2.5-7B-Instruct, Qwen3-8B, and Qwen3-30B-A3B—across six benchmarks: AlpacaEval 2.0, Arena-Hard, WritingBench, EQ-Bench 3, GPQA Diamond, and HealthBench. It outperformed the strongest static aggregation baseline in all 18 model–benchmark comparisons, with average gains up to $y_j$8 points. Rubric-level analysis grouped criteria into hard, medium, and easy buckets; the method improved hard-criterion mean scores and pass rates across all backbones and domains, maintained comparable performance on easy criteria, and produced shorter average rollouts. Mechanism analysis showed that it shifts weight mass toward high-headroom criteria, while ablations showed that replacing frontier weighting with uniform averaging degrades performance and that freezing the scalarizer makes performance almost revert to static aggregation. This pattern supports the interpretation of reward-balanced batching as online, saturation-aware reallocation of optimization effort (Huang et al., 26 May 2026).

4. Conservation-based reward balancing in bagged-reward reinforcement learning

A distinct but closely related interpretation appears in reinforcement learning from bagged reward. Here the agent does not observe per-step rewards $y_j$9, but only bag-level rewards attached to contiguous subsequences of a trajectory. A bagged reward MDP is defined as

$$\mathbf{s}_{i,j} = J(x,y_i,y_j;\mathcal{C}) \in [0,S_{\max}]^K,$$0

and the objective is to maximize expected sum of bag rewards over the bag partition of each trajectory. The reward-balanced batching perspective is explicit: each bag is a mini-batch of transitions sharing a single scalar reward, and the problem is to allocate or redistribute that reward across the instances in the bag.

The conservation constraint is

$$\mathbf{s}_{i,j} = J(x,y_i,y_j;\mathcal{C}) \in [0,S_{\max}]^K,$$1

with ideal equality. The paper trains a reward model by minimizing

$$\mathbf{s}_{i,j} = J(x,y_i,y_j;\mathcal{C}) \in [0,S_{\max}]^K,$$2

together with a next-state prediction loss

$$\mathbf{s}_{i,j} = J(x,y_i,y_j;\mathcal{C}) \in [0,S_{\max}]^K,$$3

The resulting Reward Bag Transformer uses a causal Transformer backbone to encode the state–action sequence and adds a bidirectional self-attention layer on top to predict per-step rewards: $$\mathbf{s}_{i,j} = J(x,y_i,y_j;\mathcal{C}) \in [0,S_{\max}]^K,$$4 where the attention weights $$\mathbf{s}_{i,j} = J(x,y_i,y_j;\mathcal{C}) \in [0,S_{\max}]^K,$$5 depend on queries and keys from the full sequence. The bidirectional layer is central because the observable reward is non-Markovian: reward at time $$\mathbf{s}_{i,j} = J(x,y_i,y_j;\mathcal{C}) \in [0,S_{\max}]^K,$$6 may depend on how the trajectory evolves later within the bag.

The theoretical result is that BRMDPs preserve optimal policies of the underlying MDP whenever bag rewards are true sums of per-step rewards: $$\mathbf{s}_{i,j} = J(x,y_i,y_j;\mathcal{C}) \in [0,S_{\max}]^K,$$7 Thus any redistributed per-step signal whose trajectory sum matches the bag sum is reward-equivalent at the level of optimal control; what differs across methods is learning efficiency. In this setting, reward-balanced batching means conservation-preserving, context-aware credit assignment within a batch of transitions. Experiments on Ant-v2, Hopper-v2, HalfCheetah-v2, and Walker2d-v2 with fixed-length, variable-length, and arbitrary bags showed that the Reward Bag Transformer consistently outperformed SAC with bag-end rewards, IRCR, RRD, LIRPG, and HC; in the case study on Hopper-v2 with bag length $$\mathbf{s}_{i,j} = J(x,y_i,y_j;\mathcal{C}) \in [0,S_{\max}]^K,$$8, the predicted reward time series tracked the shape and peaks of the true reward closely, and ablations showed that removing bidirectional attention or next-state prediction significantly degraded performance (Tang et al., 2024).

5. Batched exploration, experiment design, and regret balancing

In batched bandits and adaptive experimentation, reward-balanced batching usually means choosing when to update and how much data to accumulate per batch so that regret or statistical power is well distributed across the horizon. Batched Thompson Sampling provides an anytime example: the agent updates only at batch boundaries, with arm-specific cycle limits

$$\mathbf{s}_{i,j} = J(x,y_i,y_j;\mathcal{C}) \in [0,S_{\max}]^K,$$9

and ends a batch when some arm reaches its limit. This yields problem-dependent regret of order $\mathbf{w}_{\text{base}}$0, minimax regret of order $\mathbf{w}_{\text{base}}$1, $\mathbf{w}_{\text{base}}$2 batches in the worst case, and an instance-dependent expected batch count of order $\mathbf{w}_{\text{base}}$3, so batch growth is tied to posterior evidence rather than elapsed time (Kalkanli et al., 2021).

High-dimensional contextual bandits make the balancing principle explicit in the batch grid itself. For sparse and low-rank reward models, the schedule

$\mathbf{w}_{\text{base}}$4

is chosen so that per-batch regret is approximately $\mathbf{w}_{\text{base}}$5, leading to $\mathbf{w}_{\text{base}}$6 regret in the sparse case and $\mathbf{w}_{\text{base}}$7 in the low-rank case using only $\mathbf{w}_{\text{base}}$8 batches (Fan et al., 2023). In nonparametric contextual bandits, the same idea appears as dynamic binning aligned with batch size; the minimax rate with $\mathbf{w}_{\text{base}}$9 batches is

$$S_{i,j}(\mathbf{w}) = \langle \mathbf{w}, \mathbf{s}_{i,j}\rangle = \sum_{k=1}^K w^{(k)} s_{i,j}^{(k)}.$$0

and nearly constant update complexity, $$S_{i,j}(\mathbf{w}) = \langle \mathbf{w}, \mathbf{s}_{i,j}\rangle = \sum_{k=1}^K w^{(k)} s_{i,j}^{(k)}.$$1, suffices to match the fully online rate up to logarithmic factors (Jiang et al., 2024).

A parallel line of work treats limited-adaptivity pure exploration as a Gaussian sequential experiment. There the state is the posterior $$S_{i,j}(\mathbf{w}) = \langle \mathbf{w}, \mathbf{s}_{i,j}\rangle = \sum_{k=1}^K w^{(k)} s_{i,j}^{(k)}.$$2 over average rewards, and Residual Horizon Optimization solves, at each epoch, a planning problem over the remaining batch budget: $$S_{i,j}(\mathbf{w}) = \langle \mathbf{w}, \mathbf{s}_{i,j}\rangle = \sum_{k=1}^K w^{(k)} s_{i,j}^{(k)}.$$3 This directly balances posterior mean against posterior variance and uses only aggregate Gaussian approximations, not a full reward-distribution model (Che et al., 2023).

Neural contextual bandits exhibit the same tradeoff in function-approximation form. BatchNeuralUCB obtains fixed-batch regret

$$S_{i,j}(\mathbf{w}) = \langle \mathbf{w}, \mathbf{s}_{i,j}\rangle = \sum_{k=1}^K w^{(k)} s_{i,j}^{(k)}.$$4

so $$S_{i,j}(\mathbf{w}) = \langle \mathbf{w}, \mathbf{s}_{i,j}\rangle = \sum_{k=1}^K w^{(k)} s_{i,j}^{(k)}.$$5 is sufficient for matching sequential order, while determinant-triggered adaptive batching

$$S_{i,j}(\mathbf{w}) = \langle \mathbf{w}, \mathbf{s}_{i,j}\rangle = \sum_{k=1}^K w^{(k)} s_{i,j}^{(k)}.$$6

reduces the required number of policy updates to $$S_{i,j}(\mathbf{w}) = \langle \mathbf{w}, \mathbf{s}_{i,j}\rangle = \sum_{k=1}^K w^{(k)} s_{i,j}^{(k)}.$$7 while retaining the same regret order (Gu et al., 2021). Under heavy-tailed rewards, robust batched bandits replace empirical means by median-of-means estimators and reveal a distinctive asymmetry: heavier-tailed rewards require a smaller number of batches to achieve near-optimal regret in the instance-independent and Lipschitz settings, whereas the number of batches needed for near-optimal instance-dependent regret is invariant to tail heaviness (Guo et al., 4 Oct 2025).

6. System-level interpretations, adaptive control, and open issues

A broader interpretation of reward-balanced batching arises when a batch mediates between conflicting service objectives rather than between multiple reward criteria. In dynamic batch services, the control problem is formulated as minimizing

$$S_{i,j}(\mathbf{w}) = \langle \mathbf{w}, \mathbf{s}_{i,j}\rangle = \sum_{k=1}^K w^{(k)} s_{i,j}^{(k)}.$$8

the weighted sum of average response time and average power consumption, under a batch service queue with size-dependent service times. The corresponding SMDP chooses, at each decision epoch, whether to wait or to start service with a particular batch size. An abstract overflow cost yields a finite-state approximation whose space complexity and time complexity decrease by $$S_{i,j}(\mathbf{w}) = \langle \mathbf{w}, \mathbf{s}_{i,j}\rangle = \sum_{k=1}^K w^{(k)} s_{i,j}^{(k)}.$$9 and $$\Delta_{i,j}(\mathbf{w}) = S_{i,j}(\mathbf{w}) - S_{j,i}(\mathbf{w}),$$0, respectively, in the reported setup (Xu et al., 4 Jan 2025).

Cloud inference systems instantiate a related tradeoff at finer granularity. LazyBatching abandons whole-graph batching in favor of node-level batching subject to an SLA slack constraint,

$$\Delta_{i,j}(\mathbf{w}) = S_{i,j}(\mathbf{w}) - S_{j,i}(\mathbf{w}),$$1

and reports average improvements of $$\Delta_{i,j}(\mathbf{w}) = S_{i,j}(\mathbf{w}) - S_{j,i}(\mathbf{w}),$$2, $$\Delta_{i,j}(\mathbf{w}) = S_{i,j}(\mathbf{w}) - S_{j,i}(\mathbf{w}),$$3, and $$\Delta_{i,j}(\mathbf{w}) = S_{i,j}(\mathbf{w}) - S_{j,i}(\mathbf{w}),$$4 over graph batching in response time, throughput, and SLA satisfaction, respectively (Choi et al., 2020). In batched load balancing, large stale-information batches cause the “tower” phenomenon under pure Two-Choice, with gap $$\Delta_{i,j}(\mathbf{w}) = S_{i,j}(\mathbf{w}) - S_{j,i}(\mathbf{w}),$$5; mixing One-Choice and Two-Choice through the $$\Delta_{i,j}(\mathbf{w}) = S_{i,j}(\mathbf{w}) - S_{j,i}(\mathbf{w}),$$6-process, with appropriately chosen $$\Delta_{i,j}(\mathbf{w}) = S_{i,j}(\mathbf{w}) - S_{j,i}(\mathbf{w}),$$7, reduces the gap to

$$\Delta_{i,j}(\mathbf{w}) = S_{i,j}(\mathbf{w}) - S_{j,i}(\mathbf{w}),$$8

which matches a lower bound for any batched allocation process (Los et al., 2023). The earlier generalized analysis had already shown gap bounds $$\Delta_{i,j}(\mathbf{w}) = S_{i,j}(\mathbf{w}) - S_{j,i}(\mathbf{w}),$$9 for arbitrary $\Phi_\tau$0 and $\Phi_\tau$1 for $\Phi_\tau$2 across a broad class of processes, including weighted jobs and one-bit-feedback schemes (Los et al., 2022).

Lookahead-based RL provides a final control-theoretic variant. With $\Phi_\tau$3-step realized lookahead information, fixed batching policies can be exponentially suboptimal, whereas adaptive batching policies choose a state-dependent horizon $\Phi_\tau$4 and a within-batch deterministic control $\Phi_\tau$5. Their optimal Bellman equation is

$\Phi_\tau$6

which makes the batch length itself a value-maximizing control variable (Merlis, 15 Jan 2026).

A common misconception is that reward-balanced batching is synonymous with uniform weighting or with larger batches. The literature points in the opposite direction. Uniform criterion weights can entrench reward polarization in rubric-based RL; uniform redistribution inside a bag is a weak baseline for bagged reward RL; pure Two-Choice can be inferior to mixed strategies in large stale batches; and fixed, state-agnostic batching can be exponentially suboptimal when the value of lookahead is localized. The central lesson is more specific: batching is most effective when the batch boundary, the within-batch aggregation rule, and the reweighting mechanism are matched to the heterogeneity of the signal being aggregated.