Single-task RL agents are typically trained under a fixed reward function, which limits their robustness to reward misspecification and their ability to adapt to changing preferences. We introduce Reward-Conditioned Reinforcement Learning (RCRL), an off-policy method that conditions agents on reward parameterizations while collecting experience under a single nominal objective. By recomputing counterfactual rewards from shared replay data, RCRL exposes the agent to multiple reward objectives without additional environment interaction, connecting single-task RL with ideas from multi-objective and multi-task learning. Across single-task, multi-task, and vision-based benchmarks, RCRL improves sample efficiency under the nominal reward parameterization, enables efficient adaptation to new parameterizations, and supports zero-shot behavioral adjustment at deployment. Our results show that RCRL provides a scalable mechanism for learning robust, steerable policies without sacrificing the simplicity of single-task training.
Figures & tables
Figure 1: Results summary. When evaluated under the nominal rewards, RCRL substantially improves the final-step performance in both single and multi-task settings, with smaller gains in the vision-based benchmark (Figure 3 ). Furthermore, RCRL enables zero-shot transfer and policy steerability that is absent in state-of-the-art single-task RL (Figures 4 and 5 ).
Figure 2: Overview of RCRL. ( A ) During interaction, actions are sampled from the policy conditioned on the nominal reward parameterization ψ⋆ , with reward components stored in the buffer. ( B ) During training, each transition is paired with an independently sampled reward parameterization ψ∼PΨ to compute the scalar reward rψ . ( C ) The backbone training algorithm is then applied with the agent conditioned on ψ , by concatenating its representation to the environment state.
Figure 3: Sample efficiency of RCRL when evaluated under the nominal reward. Results for single-task (top row) and multi-task (bottom row) benchmarks. In the single-task RCRL uses PRC, and in the multi-task it uses ARC. Across both regimes, RCRL improves efficiency when compared to the baseline algorithms. We detail these experiments in Appendix E
Figure 4: Transfer with RCRL. We show zero-shot and finetuning performance. For the best-performing source-target task pairs, RCRL attains up to 40% of optimal performance without any finetuning, and up to 90% after 250 k environment steps. We also show heatmaps illustrating performance for all task pairs after 250k steps. The middle panel compares SimbaV2+RCRL to SimbaV2 finetuning, while the right panel compares SimbaV2+RCRL finetuning to training SimbaV2 from scratch on the target task. The results show synergies for some task pairs, alongside pairs where transfer is less effective.
Figure 5: Zero-shot transfer. Alternative rewards promote different behaviors: running speed for cheetah , standing height for hopper , and action penalty for humanoid . We compare a vanilla single-task SimbaV2 , SimbaV2+RCRL , and a full multi-task BRC agent that trains and explores under multiple reward functions. The top row shows behavioral metrics as the policy is conditioned on different parameterizations, while the bottom row shows the corresponding returns. While the vanilla agent cannot adjust its behavior without retraining, RCRL achieves behavior modulation comparable to full multi-task learning, despite learning alternative objectives fully off-policy and without collecting data under those rewards.
Figure 6: Ablation studies. ( Left ) Comparison of PRC and ARC across single and multi-task (denoted with *) benchmarks. We find that both strategies yield comparable gains. ( Middle ) Effect of conditioning the agent on the parameterization ψ in 14 HumanoidBench tasks. Removing conditioning degrades performance, with a larger impact for ARC due to greater reward diversity. ( Right ) Conditioning via input and via output (i.e. multi-head architecture ( Yu et al., 2019 ) ) on 14 HumanoidBench tasks. While multi-head model does not scale to arbitrary number of parameterizations, it also yields highly competitive performance.
Figure 7: Ablation studies. Experiments ran on 23 tasks from from HB, DMC, and Gym. ( Left ) Comparison of reward conditioning during learning and exploration. A variant that both explores and learns under sampled parameterizations also outperforms the baseline, showing that parameterizations are beneficial not only during learning, but also during exploration. ( Middle ) Sensitivity to the conditioning probability α , which controls the fraction of updates performed under the nominal reward. Performance varies smoothly, indicating robustness to this hyperparameter. ( Right ) We evaluate both SimbaV2 and SimbaV2+RCRL with small and large networks. While RCRL performs similarly to the baseline at small scale, it gains more from increasing model capacity, suggesting that its supervision signal scales better than standard single-task supervision.
Appendix figures & tables16 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Code example. We present the code used to sample perturbations vector Δ , as well as a OpenAI Gym wrapper used to calculate alternative reward functions in our single-task experiments.
Figure 9: Code example. We present the code used to calculate auxiliary task rewards for the DeepMind Control dog tasks. We use this wrapper in our multi-task experiments.
DMC Dogs
HB
HB Hard
stand
crawl
crawl
walk
stand
stand
trot
walk
walk
run
run
run
( stair , slide )
( stair , slide )
pole
pole
Appendix
Table 1: Tasks used for auxiliary reward conditioning. Tasks grouped in parentheses share identical reward functions and are treated as a single reward specification.
Figure 10: Code example. We present the code used to calculate auxiliary task rewards in the zero-shot policy adjustment experiments.
Table 12
Figure 11: Ablation studies. ( Left ) Comparison of input representations for PRC in the single-task HumanoidBench setting. We compare conditioning via a learned embedding module, optimized through the critic loss, against directly concatenating the perturbation vector Δ to the state on 14 tasks. ( Middle ) Comparison of perturbation sampling distributions for PRC on the same 14 HumanoidBench tasks. We compare the default log-uniform distribution described in Section against a log-normal distribution centered around the nominal parameterization. ( Right ) We vary the spread of the log-uniform distribution used to sample reward perturbations spread∈{2,4,8,16} while keeping all other hyperparameters fixed.
Benchmark
Task
Random Score
Optimal Score
DeepMind Control
dog-stand
0.000
1000
DeepMind Control
dog-walk
0.000
1000
DeepMind Control
dog-trot
0.000
1000
DeepMind Control
dog-run
0.000
1000
OpenAI Gym
Ant-v4
-70.288
3942
OpenAI Gym
HalfCheetah-v4
-289.415
10574
Appendix
Table 5: Random and optimal scores used for normalization.
Figure 12: Understanding RCRL. We decompose the gains of multi-task training on the DMC-Dogs benchmark into broader state-action coverage and additional reward supervision. ST denotes single-task training, MT denotes full multi-task training, ST+RCRL denotes single-task data collection with counterfactual reward supervision, and ST+Expanded (s,a) denotes single-task reward training with additional trajectories collected under auxiliary tasks. MT performs best overall, while both ST+RCRL and ST+Expanded (s,a) improve over ST , suggesting that reward supervision and state-action coverage both contribute to multi-task gains. Notably, RCRL also improves easier tasks such as dog-stand , indicating that harder alternative rewards can provide useful augmentation even when the nominal task is simple.
Figure 13: Single-task DeepMind Control performance. We present mean and 95% confidence intervals calculated using bootstrapping. 10 random seeds.
Figure 14: Single-task Gym performance. We present mean and 95% confidence intervals calculated using bootstrapping. 10 random seeds.
Figure 15: Single-task HumanoidBench performance. We present mean and 95% confidence intervals calculated using bootstrapping. 10 random seeds.
Figure 16: Multi-task DeepMind Control performance. We present mean and 95% confidence intervals calculated using bootstrapping. 5 random seeds.
Figure 17: Multi-task HumanoidBench performance. We present mean and 95% confidence intervals calculated using bootstrapping. 5 random seeds.
Figure 18: Multi-task HumanoidBench Hard performance. We present mean and 95% confidence intervals calculated using bootstrapping. 5 random seeds.
Figure 19: Vision-based DeepMind Control performance. We present mean and 95% confidence intervals calculated using bootstrapping. 5 random seeds.
Figure 20: Transfer finetuning performance. We present mean and 95% confidence intervals calculated using bootstrapping. 5 random seeds.
Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.
Offline zero-shot reinforcement learning (RL) aims to learn agents that optimize unseen reward functions without additional environment interaction. The standard approach to this problem trains task-conditioned policies by sampling task vectors that define linear reward functions over learned state representations. In most existing algorithms, these task vectors are randomly sampled, implicitly assuming this adequately captures the structure of the task space. We argue that doing so leads to suboptimal zero-shot generalization. To address this limitation, we propose extracting task vectors directly from the offline dataset and using them to define the task distribution used for policy training. We introduce a simple and general reward function extraction procedure that integrates into existing offline zero-shot RL algorithms. Across multiple benchmark environments and baselines, our approach improves zero-shot performance by an average of 20%, highlighting the importance of principled task sampling in offline zero-shot RL.
Nazim Bendib, Nicolas Perrin-Gilbert, Olivier Sigaud
Institut des Systèmes Intelligents et de Robotique, ISIR, Sorbonne Université, Paris, France.
Contextual Reinforcement Learning (CRL) seeks to generalize classical RL by maximizing task coverage across a context space of related tasks. While prior works often train from scratch and rely on either multi-task learning for a single policy or strategically training multiple policies, we advocate for a unified alternative: pretraining a single policy with good initial performance, followed by fine-tuning multiple policies for task specialization. This new paradigm, however, introduces unique challenges, such as heterogeneous marginal returns and sample inefficiency. This raises a critical research question: given a pretrained policy and a constrained budget, how much fine-tuning should each task region receive to enable sample-efficient CRL? To this end, we propose Task Specialization Fine-Tuning (TSFT), an online framework that predicts fine-tuning performance with a simple parametric model and exactly solves the resulting discrete budget allocation problem via integer linear programming. Extensive experiments across diverse decision domains, including combinatorial optimization, continuous control, and LLM fine-tuning, demonstrate that TSFT significantly outperforms baselines in task coverage and approaches oracle performance. Our work charts a new direction for model-based CRL, aligning with the modern pretrain-finetune era.
In-context reinforcement learning (ICRL) has emerged as an effective paradigm for test time adaptation to unseen tasks without parameter updates. However, existing ICRL methods can exhibit brittle and unstable adaptions, and the mechanisms underlying such adaptions remain poorly understood. We provide a policy-gradient view of ICRL and argue that relying on trajectory-level return feedback can lead to high-variance updates when test time rollout is limited. Motivated by the actor--critic principle, we propose Context Value Informed ICRL (CV-ICRL), which equips the in-context policy with an explicit context value: the expected discounted return conditioned on the current state and accumulated interaction context. CV-ICRL trains a value head and writes its prediction back into the context as a value token, enabling TD-style return targets for lower-variance policy updates. Experiments on the Dark Room, Minigrid, and Procgen testbeds show that CV-ICRL substantially improves the stability of test time adaptation and achieves higher returns across tasks and environments. Code and data are available at https://github.com/Bluixe/CV-ICRL .