Title: Towards Robust Zero-Shot Reinforcement Learning
Venue: NeurIPS 2025
Accepted: yes
Keywords: Reinforcement Learning

Zero-shot Reinforcement Learning

Offline Reinforcement Learning
Abstract: The recent development of zero-shot reinforcement learning (RL) has opened a new avenue for learning pre-trained generalist policies that can adapt to arbitrary new tasks in a zero-shot manner. While the popular Forward-Backward representations (FB) and related methods have shown promise in zero-shot RL, we empirically found that their modeling lacks expressivity and that extrapolation errors caused by out-of-distribution (OOD) actions during offline learning sometimes lead to biased representations, ultimately resulting in suboptimal performance. To address these issues, we propose Behavior-REgularizEd Zero-shot RL with Expressivity enhancement (BREEZE), an upgraded FB-based framework that simultaneously enhances learning stability, policy extraction capability, and representation learning quality. BREEZE introduces behavioral regularization in zero-shot RL policy learning, transforming policy optimization into a stable in-sample learning paradigm. Additionally, BREEZE extracts the policy using a task-conditioned diffusion model, enabling the generation of high-quality and multimodal action distributions in zero-shot RL settings. Moreover, BREEZE employs expressive attention-based architectures for representation modeling to capture the complex relationships between environmental dynamics. Extensive experiments on ExORL and D4RL Kitchen demonstrate that BREEZE achieves the best or near-the-best performance while exhibiting superior robustness compared to prior offline zero-shot RL methods. The official implementation is available at: https://github.com/Whiterrrrr/BREEZE.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces several improvements to the FB algorithm to improve zero-shot RL. Specifically, they use attention for F and B, they use an additional value function to stabilize training, they use diffusion for the policy, and they use an IQL style constrained update for the value function. They demonstrate that basic FB does not scale with additional parameters. Furthermore, their algorithm outperforms FB and slightly outperforms HILP in terms of downstream return.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1) The visuals look good and are useful.\n\n2) Their explanations of *what* they did is clear.\n\n3) Their experiments are thorough - many environments, seeds, reasonable baselines. \n\nWeaknesses:\n1) I think this paper is missing the *why*. Why is attention useful for F and B? Does some other aspect of attention explain the improved performance? It could be layer norm or residual connections explaining this difference. It is unclear to me why attention itself is useful here, and the authors do not motivate this other than to point to empirical results. \n\n2) In this setting, regulating towards the exploratory policy may not be a good idea. The exploratory policy is likely not optimal for a downstream reward function. Therefore we want to deviate from the exploratory policy in order to maximize reward. While this does run the risk of OOD model updates, it seems necessary to transfer. In contrast to normal offline RL, we are trying to achieve a different objective from the offline data's policy. Why is staying close to the exploratory policy a good design choice?\n\n3) The connections to negative values out of FB is not fully expanded on or solved directly. They mention this is a problem for FB, and that their method outputs negative values less often, but still does. If the only argument is that negative values are a symptom of approximation error, and therefore they have less approximation error, I think this argument is only somewhat convincing. I also think the discussion of large magnitude values being an issue for FB may not be correct. \n\n\nMinor: \n\nNotation at line 85 is weird. You use big S_+ in equation 1, but use little s_+ in S in line 86. \n\nTypo Line 88\n\nTypo line 109\n\nTypo figure 2\n\nLine 170: by solving *the* following dataset-constrained problem\n\nLine 204: typo\n\nLine 216 typo\n\nQUESTIONS:\n1) In section 3, you mention that FB distributions \"exhibit large absolute values\". But the measure is effectively a sum of distributions, and if the transition function is deterministic, the measure for the next state should be a dirac, ie infinity. In other words, how do you know that these values are too large? Infinity is potentially a correct value, although I agree this causes substantial problems in practice. \n2) Equation 4 and 5: By incorporating E(BB^T), you are effectively projecting the reward function corresponding to z onto the basis B via least squares. Therefore, B is effectively learning a linear span of reward functions. Please comment on this design decision. I think this will change what B represents. \n3) Equation 6: z is an input to V_pi_z on the left, but is maxxed over on the right. I think you mean to only max over a in A? \n4) Equation 6: How do you know mu(a|s)? Are you assuming access to the exploratory policy? If not, are you estimating this? You cannot use the empirical distribution directly because you will never visit the same state twice in a continuous setting. \n5) Why do you use attention based architecture when the inputs are only 1 or 2 tokens? Especially in the case of 1 token, this is not intuitive because there is nothing else to attend to. How do you know which part of the transformer is necessary here? Maybe only layer norms are needed, dropout, residual connections etc. It is not clear why attention is useful in this case. \n6) For figure 11,12,13, why does HILP only train for half the time? \n7) From the training curves, this algorithm seems to converge asymptotically to a similar place as other algorithms. I see the primary advantage here in terms of convergence speed. However, there are many expensive components added, including attention, diffusion, and the additional learned value function. How does your approach compare to others in terms of compute speed? \n8) What happens if tau=1? From line 174 it seems like it should cause problems, but the trend of the ablation suggests it may perform well.\n\nLIMITATIONS:\nLargely this work needs ablations to help explain the \"why\" of this paper. \n1) This work needs ablations on what about attention actually improves things over vanilla NNs\n2) This work needs an ablation showing the compute time relative to prior works. Many of the components are expensive: Attention, diffusion, extra value function. How do the training curves compare when the x axis is training time rather than update steps? It is possible the improved performance is due to using significantly more compute per gradient step. \n3) This work needs an ablation on the IQL value loss. What happens if you train the value function using a typical TD-type loss instead? Is the IQL loss what improves performance, or the use of an additional value function?", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes several improvements to the Forward-Backward Representations (FB) framework for zero-shot RL. The FB framework represents the successor measure $M^{\\pi_z}(s_0, a_0, s^+)$ using the dot product of a forward representation $F(s_0, a_0, z)$ encoding the future state occupancy given state-action pair and task policy $\\pi_z$, and a backward representation $B(s^+)$ encoding the state distribution leading up to $s^+$. A $z$-conditioned policy is jointly trained to decode the optimal policies for each task vector $z$. The paper is motivated by the observation that FB often learns heavily distorted $M$ and $Q$. This is be attributed to a few reasons: \n1. The FB objective does not enforce proper scaling, \n2. The policy class is not expressive enough to encode all possible optimal policies\n3. No pessimism for OOD states. \n\nTo address these, the paper proposes the following improvements:\n\n1. Regularize the $F$ network using dataset samples by replacing value target in $F$ network objective with a separate value network (as opposed to bootstraping from $F$ network). The value network is trained with expectile loss on the dataset, ensuring proper normalization while being pessimistic OOD.\n2. Use a self-attention architecture to parameterize the $F$ and $B$ networks for better expressivity.\n3. Use a diffusion model to parameterize the policy. The model is trained via a weighted-regression objective which correspond to learning the optimal policy for each $z$.\n\nThese contributions results in state-of-the-art performance on multi-task RL problems on ExoRL and D4RL tasks, outperforming baselines such as successor features, FB, and HILP. Overall, the methods proposed in this paper greatly improves the practicality of the FB framework.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n1. The paper is well-motivated from an empirical analysis of the failure modes of the FB algorithm.\n2. The proposed solutions are sensible and apply well-established machineries in the literature.\n3. The claims made in this paper are supported by extensive experiments on ExoRL and D4RL domains. Each individual component is justified via ablation experiments. \n4. The derivations of the theorems and lemmas are sound.\n\n**Weaknesses**\n1. The paper can feel like an incremental contribution to / modernization of the FB framework.\n2. A lot of extra machinaries and compute are required to get marginal improvements over strongest baseline (HILP). According to the ablations in Section 5.2, removing a single component can cause performance to drop below baseline.\n\nQUESTIONS:\n1. How should I interpret the $M$ and $Q$ plots in Figure 1? Why is positive values more favorable? Are there ground truth $M$ and $Q$ to compare to?\n2. FB has a regularizer to ensure proper scaling of $F$ and $B$ ([1] Appendix. B, orthonormalization loss). Does this address the $M$ and $Q$ scaling problem?\n3. For policy extraction, can you do guided diffusion on the behavior policy by adding $\\alpha(F(s,a,z)^{\\top}z - V_{\\pi_z}(s,z))$ to the noise prediction at each timestep?\n4. The statement of Lemma 1 in Appendix A.2 refers to a missing citation [42].\n\n[1] Ahmed Touati, Jérémy Rapin, Yann Ollivier. Does Zero-Shot Reinforcement Learning Exist? ICLR 2023.\n\nLIMITATIONS:\nThe authors addressed the limitations of the paper in Appendix D.4, stating that the proposed method introduce extra compute compared to the standard FB framework.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an extension to Forward-Backward (FB) models in offline reinforcement learning by incorporating an improved out-of-distribution regularization scheme alongside enhanced attention-based network architectures. The motivation stems from shortcomings identified in prior FB approaches. The proposed method is evaluated on a wide range of offline RL tasks, including different datasets and domains, and is benchmarked against several state-of-the-art baselines. Results indicate that the method achieves overall improved performance and robustness, supported by extensive ablations and visualizations.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n- Realistic Contribution: The paper makes a practically grounded contribution, avoiding overclaiming or speculative framing.\n- Clear Motivation and Formal Soundness: The approach is well-motivated by formal considerations of current model shortcomings, and the proposed methodology is solidly defined.\n- Comprehensive Evaluation: The experimental setup is broad, covering a range of tasks, datasets, and baselines. Ablation- and hyperparameter studies provide insight into the contribution of each component. The illustrative visualizations and provided results strongly support the empirical claims.\n\n## Weaknesses\n\n- Clarity and Writing: Overall readability suffers due to long and dense sentences. Splitting them for clarity would improve comprehension significantly.\n- Preliminaries and Task Formalization: Including the discount factor with the reward free MDP seems unnecessary. Also, I am missing a thorogh formalization of the considered task.\n- Ambiguity in Zero-shot Framing: The terminology around \"zero-shot\" learning is ambiguous (e.g., \"using only a small set of reward-labeled samples\" could also be conceived as few-shot). A central definition would help.\n- Missing Limitations and Assumptions: The paper lacks a clear discussion of limitations, such as assumptions on the availability and quality of offline data, or the implications of the KL constraint.\n- Related Work: Could be improved by better positioning the contribution relative to prior approaches. Also, the improved representation modeling proposed could be considered related to the partial observability considered in [1].\n- Computational Details: A discussion on model complexity, runtime, and training resource requirements is currently missing. These are important for understanding practical applicability.\n\n[1] Altmann et al., \"CROP: Towards Distributional-Shift Robust Reinforcement Learning using Compact Reshaped Observation Processing\", in Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence, 2023.\n\nQUESTIONS:\nWhat exactly constitutes the “ground truth” in Figure 1? The figure’s visual semantics do not clearly reflect the described differences in the text. Additionally, BREEZE also appears to show high values, contradicting the written interpretation.\n\nCould the authors elaborate on what “zero-shot” means in their context? The usage could be interpreted as few-shot instead, especially since some reward-labeled data is used. A precise definition would help.\n\nCould the authors discuss potential failure cases, data conditions, or general assumptions limiting their approach? This would strengthen the claims of robustness and highlight future directions for improvement.\n\nLIMITATIONS:\nNo. Limitations are not discussed. Suggested areas to address include:\n- Dependence on large, high-quality offline datasets.\n- Computational cost and scalability.\n- Potential brittleness to domain shift beyond training datasets.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nBREEZE is a behaviour-regularised, expressivity-enhanced Forward-Backward (FB) framework for zero-shot reinforcement learning. It (i) adds a behaviour-constrained expectile-value objective that keeps policy learning in-sample to curb OOD action extrapolation, (ii) replaces shallow MLP encoders with self-attention forward/backward networks for richer successor-measure estimates, and (iii) extracts actions with a guided task-conditioned diffusion policy to handle multi-modal action distributions. Across 12 ExORL domains and four long-horizon Franka-Kitchen tasks, BREEZE sets new state-of-the-art IQM returns and shows markedly smoother learning than prior FB variants .\n\nSTRENGTHS AND WEAKNESSES:\n- Principled behaviour regularisation: The KL-bounded, expectile-value loss forces learning to stay on the dataset support, directly attacking OOD action extrapolation .\n- Expressive FB encoders: Two-token self-attention forward and transformer backward nets capture complex state–task relations far better than widened MLPs, lifting FB performance noticeably (Fig. 2) .\n- Diffusion actor models multi-modal actions: A task-conditioned diffusion policy with value guidance generates diverse candidate actions and picks the top-$K$ by $Q$, yielding higher returns on high-variance domains .\n- Strong, broad empirical gains: BREEZE tops 10/12 ExORL tasks and all Kitchen benchmarks in IQM return and variance, clearly outperforming FB/MCFB/VCFB and HILP baselines .\n- Informative diagnostics & ablations: Distribution plots reveal scale errors in previous FB methods, while component ablations isolate the impact of FB enhancement and diffusion actor (Table 2, Fig. 6) .\n\nQUESTIONS:\n- Compute footprint & efficiency – Please report training time, GPU type/count, and peak memory for ExORL/Kitchen runs; show BREEZE stays within 1.5 × wall-clock of FB or explain the trade-off.\n- Generality beyond MuJoCo/Kitchen – Evaluate on at least one pixel-based suite (e.g., DMControl-Pixels) or Atari-URLB. Comparable gains would strengthen significance; failure analysis would also help.\n- Ablate diffusion on Kitchen – Add a Kitchen variant using a Gaussian actor to clarify when diffusion is essential.\n- Failure-case deep dive – For tasks with variance spikes, supply diagnostics linking instability to either critic bias or diffusion sampling.\n\nLIMITATIONS:\n- Limited domain diversity: All experiments are MuJoCo or Franka-Kitchen; no pixel-based control or real-robot data are tested, so generality is uncertain .\n- parameter sensitivity: Expectile $τ$ and guidance weight $α$ demand careful tuning; Figure 6 shows sharp performance cliffs with poor choices .\n- Sparse analysis of failure cases: Experiment results show variance spikes, but the paper does not dissect whether these stem from value mis-estimates or diffusion sampling.\n- Missing comparison to function-encoder ZSRL lines: Recent methods like FROE or F-Reward are omitted, so progress relative to non-FB approaches is unclear .", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-07"}