Bigger, Regularized, Categorical: High-Capacity Value Functions are Efficient Multi-Task Learners

1 Introduction

• Naïvely scaling SAC to multi-task learning decreases aggregate performance, whereas BRC improves both single-task and multi-task results.Figure 1 compares naïve multi-task SAC with BRC and prior approaches, including optimized single-task learners.Source: Figure 1 · p. 1
• BRC combines high-capacity value models, cross-entropy loss, and learned task embeddings to scale online multi-task RL.These choices target gradient-magnitude variation and conflicts that make naïve multi-task value learning brittle.Source: §1 · p. 2, Figure 3 · p. 3
• Multi-task BRC improves final performance and sample efficiency over single-task learning while using fewer gradient updates.The introduction attributes this result to broader task data and reports the comparison in Figure 1.Source: Figure 1 · p. 1, §1 · p. 2
• A pretrained multi-task BRC model transfers to new tasks with improved sample efficiency compared with learning from scratch.The transfer result is presented in Figure 2.Source: §1 · p. 2, Figure 2 · p. 2
• Cross-entropy with return normalization stabilizes multi-task learning by reducing variability caused by different reward magnitudes.Figure 3 contrasts MSE, normalized MSE, and normalized cross-entropy on HB-MEDIUM.Source: §1 · p. 2, Figure 3 · p. 3

2 Related Work

• Earlier studies found that naïvely increasing model capacity can cause divergence under temporal-difference training.Layer normalization, feature normalization, and classification losses were explored as stabilizing mechanisms.Source: §2 · p. 3
• Multi-task RL has been challenged by heterogeneous reward scales and conflicting gradients, motivating normalization, projection, and policy-distillation methods.These approaches address instability without the paper’s combined capacity-and-cross-entropy strategy.Source: §2 · p. 3
• Generalist agents are often trained by behavioral cloning on expert demonstrations because online multi-task RL is difficult.Prior work also uses pretrained vision and language backbones and fine-tuning for transfer.Source: §2 · p. 3

Cross-entropy loss

• Cross-entropy training reduces reward-scale-driven differences in TD errors and gradient magnitudes that destabilize multi-task learning.MSE is sensitive to predicted-value scale and outliers, while higher-reward tasks can otherwise produce stronger learning signals.Source: §3 · pp. 3–4

Task embeddings

• Task embeddings outperform separate heads across all considered model scales and learn representations that cluster tasks with similar embodiments.The embeddings are learned online and reveal structure in the underlying dynamics.Source: Figure 5 · p. 4, Figure 7 · p. 5

Scaled Q-value model

• BroNet (CE) improves performance as model capacity grows in both single-task and multi-task HB-MEDIUM, unlike vanilla architectures whose performance degrades with size.The proposed architecture combines BroNet with cross-entropy loss and achieves the best performance in both settings.Source: §3 · p. 4, Figure 4 · p. 4

Task interference

• Increasing parameter count reduces the rate of gradient conflicts, indicating that higher-capacity models can mitigate a key source of multi-task interference.The comparison is made among single-task, separate-head, and task-embedding variants on HB-MEDIUM.Source: §3 · p. 4, Figure 5 · p. 4

Component analysis

• The scaled Q-value model is the most influential BRC component, accounting for 62% of the measured improvement, while removing any component decreases performance.Figure 6 evaluates combinations of scaled Q-values, cross-entropy loss, and task embeddings using Shapley values.Source: Figure 6 · p. 5, §3 · p. 5

BRC design

• BRC combines scaled residual Q-networks, cross-entropy critic training, and learnable task embeddings for scalable online multi-task RL.The approach is applied with SAC and DrQ-ϵ to continuous and discrete control.Source: §3 · p. 5

Experimental questions

• The experiments test whether capacity scaling improves multi-task performance, whether task count improves learning, and whether learned representations transfer efficiently to new tasks.The study also examines how model capacity affects task scaling and transfer.Source: §4 · pp. 5–6

Benchmarks

• The evaluation covers 283 tasks across DeepMind Control, MetaWorld, HumanoidBench, Atari, and ShadowHand, including sparse-reward, high-degree-of-freedom, and vision-based control.The experiments use SAC for continuous control and DrQ-ϵ for discrete vision-based control with shared hyperparameters across tasks.Source: §4 · p. 6, Figure 8 · p. 6

Transfer learning

• Transfer evaluation uses weight initialization, embedding tuning from fewer than a thousand target-task transitions, and data transfer from a pretrained experience buffer.Transfer tasks are excluded from multi-task pretraining to preserve a train-test split analogous to supervised learning.Source: §4 · p. 6

Evaluation

• Performance is normalized to [0, 1] using domain-specific return or success-score conventions, with aggregate results reported as bootstrapped sample means and 95% confidence intervals.Normalization differs across DMC, MetaWorld, HumanoidBench, and Atari.Source: §4 · pp. 6–7

Scaling data and compute

• Multi-task BRC achieves a 2× sample-efficiency improvement while using 40× fewer gradient updates than state-of-the-art single-task learners on HB-HARD.Figure 9 compares proposed single- and multi-task BRC agents on sample efficiency and gradient-update requirements.Source: Figure 9 · p. 7

Efficiency of multi-task learning

• Multi-task BRC matches or surpasses strong single-task specialists while using 40× fewer gradient updates than the best single-task model.This result holds with one hyperparameter set across tasks.Source: §5 · p. 7

Scaling model capacity

• Larger models improve multi-task performance beyond 250M parameters, reduce gradient conflicts, and make Q-value scaling the largest contributor to gains.Q-value scaling accounts for more than 60% of improvements according to Shapley values.Source: §5 · p. 7

Scaling number of tasks

• Using more pretraining tasks improves both performance on shared tasks and transfer to out-of-distribution tasks, with robustness as robot embodiments increase.The transfer trend is most pronounced for the biggest evaluated models.Source: §5 · p. 7, Figure 10 · p. 8

Model transfer from multi-task RL

• Pretrained multi-task value models transfer faster and achieve higher returns than learning from scratch or initializing with the same pretraining data.Model transfer is often more practical than data transfer because pretrained weights encode the dataset without inducing distribution mismatch.Source: §5 · p. 8, Figure 11 · p. 8

Embedding-tuning transfer

• Task-embedding-only adaptation achieves over 90% success on 29 unseen shapes after selecting embeddings with 300 transitions while freezing the network.The result suggests transfer can occur through the input space without updating policy or value parameters.Source: §5 · pp. 8–9

Multi-task value models learn more efficiently with online data

• Online BRC outperforms offline BRC by nearly 20%, and the 1B model reaches offline TD-MPC2 performance with 64k transitions, almost 100× fewer samples.Offline BRC slightly underperforms TD-MPC2 on the MW+DMC benchmark.Source: §5 · p. 9, Figure 12 · p. 9

• Online temporal-difference value-function training scales beyond billion-parameter models and many tasks.Source: §6 · p. 9
• Cross-entropy-trained Q-value models conditioned on learnable task embeddings address varying gradient magnitudes and conflicting gradients.Source: §6 · p. 9
• The approach achieves state-of-the-art performance, matches single-task experts with 40× fewer gradient steps, and improves transfer sample efficiency and final performance.Source: §6 · p. 9

• The authors identify no specific societal impact beyond improving the general capacity of value-based RL agents.Source: Societal Impact · p. 17

A Limitations

• Task synergy remains unresolved: shared-embodiment tasks usually reinforce one another, whereas mixed-embodiment suites show synergy in some cases and none in others.The authors identify the conditions that create synergy as an important direction for future work.Source: A Limitations · p. 17
• When task synergy is absent, the method still learns but at a slightly slower rate, while prior approaches often failed to converge.Source: A Limitations · p. 17

Implementation limitations

• Categorical RL is challenged by multi-task reward-scale differences because its return distribution uses fixed support.Source: §C.2 · p. 18
• Per-task return normalization bounds modeled returns by Vmax and maintains similar return scales between tasks, but does not resolve the paper's open synergy question.Source: A Limitations · p. 17, §C.2 · p. 18

• Learned task embeddings are updated by backpropagating the temporal-difference loss and are constrained to have L1 norm 1.Source: §C.4 · p. 19
• The approach uses distinct observation-concatenation strategies for proprioceptive and vision-based tasks, with task embeddings also attached to the vision encoder output.Source: Figure 15 · p. 19, §C.4 · p. 19
• For proprioceptive control, the task embedding is concatenated directly with the raw state before entering the actor and critic.Source: §C.4 · p. 19
• For vision-based control, an IMPALA encoder processes RGB observations before its image features are concatenated with the task embedding and sent to the Q-network.Source: §C.4 · p. 19, Figure 15 · p. 19

Architectures and design choices

• The evaluation covers critic architectures, model widths, reward normalization, cross-entropy versus non-distributional learning, and multi-task heads on HumanoidBench-Medium.Source: §D.1 · pp. 20–21

Training comparisons

• Figure 13 compares training curves for single-task and multi-task methods on HumanoidBench-Hard using both environment steps per task and total gradient steps.Source: §D.1 · p. 21

Scaling

• Scaling experiments examine final performance across critic widths and task-subset sizes, including 3-, 6-, and 9-task HumanoidBench-Medium settings.Source: §D.1 · p. 21

Transfer

• Transfer experiments average the first 500k training steps over 15 tasks and compare scratch training with multi-task initialization and replay-buffer initialization.Source: §D.1 · p. 21

Efficiency

• On MW+DMC, the figures compare offline final performance and GPU cost, and measure environment steps needed to reach fractions of final TD-MPC2 performance at different model widths.Source: §D.1 · p. 21

Baselines

• The benchmark suite includes comparisons with SAC, Multi-Head SAC, BRO, SimBa, SimBaV2, TD-MPC2, DreamerV3, PCGrad, PaCo, and Cross-Task Policy Guidance.Source: §D.2 · pp. 21–22

Benchmark coverage

• The benchmark suite includes MetaWorld’s 50 tasks, Atari’s 26 tasks, ShadowHand’s 85 manipulation shapes, and 80 diverse MetaWorld and DeepMind Control tasks.Source: §E.1 · pp. 22–23

Humanoid benchmarks

• HumanoidBench-Medium contains 9 common-embodiment locomotion tasks with 51-dimensional states and 19-dimensional actions.Source: §E.1 · p. 23
• HumanoidBench-Hard contains 20 locomotion and manipulation tasks with varying state dimensionality, including 307-dimensional states and 61-dimensional actions.Source: §E.1 · p. 23

Other embodiments

• The suite also includes four DeepMind Control dog tasks with 223-dimensional states and 38-dimensional actions.Source: §E.1 · p. 23

Transfer splits

• Transfer evaluations leave out 20 Atari tasks, 5 MetaWorld tasks, 3 HumanoidBench-Medium tasks, and 7 HumanoidBench-Hard tasks from multi-task training.Source: §E.2 · p. 24

Evaluation protocol

• Scores are normalized to bound returns between 0 and 1 where applicable, while MetaWorld reports unnormalized success rates and Atari uses scores from Agarwal et al. [4].Source: §E.3 · p. 25
• A single hyperparameter configuration is used across all tested tasks to assess the robustness of the approach.Source: F Hyperparameters · p. 26

• The section also lists hyperparameters for DrQ+BRC and provides an Atari-100k benchmark-score table, without supplying their values here.The supplied passages identify these artifacts but do not report their contents.Source: Table 4 · p. 26, Table 5 · p. 28
• Training curves are reported for aggregate performance, HumanoidBench-Medium, HumanoidBench-Hard, DMC-Hard, MetaWorld, ShadowHand, MW+DMC, and transfer tasks.The figures use environment steps as the x-axis and benchmark-specific returns or success rates as performance measures.Source: Figure 16 · p. 26, Figure 17 · p. 27, Figure 18 · p. 27
• The reported training curves include 95% confidence intervals calculated via bootstrapping.This uncertainty reporting appears across the aggregate, benchmark-specific, mixed-benchmark, and transfer figures.Source: Figure 16 · p. 26, Figure 17 · p. 27, Figure 18 · p. 27
• The MW+DMC and transfer plots combine benchmark-specific metrics, using success rate for MetaWorld and summed episodic returns for DMC or other tasks.These mixed-benchmark plots therefore do not use one common y-axis metric across all tasks.Source: Figure 22 · p. 31, Figure 23 · p. 32
• Most testing shapes can be manipulated by the majority of pretrained embeddings, suggesting that individual policies do not overfit to particular tasks.Figure 24 evaluates pretrained embeddings on new shapes.Source: Figure 24 · p. 32