cs.LGMar 5, 2026

Reward-Conditioned Reinforcement Learning

Authors: Michal Nauman, Marek Cygan, Pieter Abbeel

Organizations: University of Warsaw · University of Warsaw, Nomagic · UC Berkeley, Amazon FAR

Abstract

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

Panel of Figure 1 — 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 ).

Panel of Figure 1 — 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

\psi^{\star}

, with reward components stored in the buffer. ( B ) During training, each transition is paired with an independently sampled reward parameterization

\psi\sim\mathcal{P}_{\Psi}

to compute the scalar reward

r_{\psi}

. ( C ) The backbone training algorithm is then applied with the agent conditioned on

\psi

, by concatenating its representation to the environment state.

Panel of Figure 3 — 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

Panel of Figure 3 — 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

Panel of Figure 4 — 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.

Panel of Figure 5 — 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.

Panel of Figure 5 — 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.

Panel of Figure 6 — 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 $\psi$ 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.

Panel of Figure 6 — 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 $\psi$ 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.

Panel of Figure 7 — 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 $\alpha$ , 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.

Panel of Figure 7 — 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 $\alpha$ , 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 \(\Delta\) , as well as a OpenAI Gym wrapper used to calculate alternative reward functions in our single-task experiments. — Figure 8: Code example. We present the code used to sample perturbations vector $\Delta$ , 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

Panel of Figure 11 — 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 $\Delta$ 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 $\text{spread}\in\{2,4,8,16\}$ while keeping all other hyperparameters fixed.

Panel of Figure 11 — 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 $\Delta$ 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 $\text{spread}\in\{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.

Panel of Figure 12 — 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.

Panel of Figure 12 — 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.

More in this topic

Reinforcement Learning

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

See all 1,661 papers

Improving Zero-Shot Offline RL via Behavioral Task Sampling
Apr 28, 2026Nazim Bendib, Nicolas Perrin-Gilbert, Olivier SigaudZero-Shot RL Offline RL
Institut des Systèmes Intelligents et de Robotique, ISIR, Sorbonne Université, Paris, France.
Task Specialization Fine-Tuning for Contextual Reinforcement Learning
Aug 17, 2026Jianan Zhou, Jung-Hoon Cho, Tianyue Zhou +5Contextual RL Reinforcement Learning
Nanyang Technological University · MIT · UIUC
Context Value Informed In-Context Reinforcement Learning
Sep 27, 2025Wenhao Zhang, Shao Zhang, Xihuai Wang +2In-Context RL Test-Time RL
Shanghai Jiao Tong University

Apr 28, 2026J/K move · Enter open · S save

Improving Zero-Shot Offline RL via Behavioral Task Sampling

Nazim Bendib, Nicolas Perrin-Gilbert, Olivier Sigaud

Institut des Systèmes Intelligents et de Robotique, ISIR, Sorbonne Université, Paris, France.

Zero-Shot RL Offline RL Multi-Task RL Reinforcement Learning Generalization in RL Goal-Conditioned RL+1

Open paper AskPDF arXiv

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.