Towards Robust Zero-Shot Reinforcement Learning
Authors: Kexin ZHENG
Lauriane Teyssier
Yinan Zheng
Yu Luo
Xianyuan Zhan
Venue: NeurIPS 2025 poster ()
Keywords: Reinforcement Learning
Zero-shot Reinforcement Learning
Offline Reinforcement Learning
Primary area: reinforcement_learning
TL;DR: BREEZE is a new zero-shot reinforcement learning framework that improves the stability, expressivity, and performance of pre-trained generalist policies, achieving state-of-the-art results on benchmark tasks.

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.

Towards Robust Zero-Shot Reinforcement Learning
Kexin Zhengl-: 1-1 Lauriane Teyssier** Yinan Zheng', Yu Luo?, Xianyuan Zhan
The Chinese University of Hong Kong TTinnguua University
Huawei Noah's Ark Lab Shanghai Artificial Intelligence Laboratory
115517372391inf cuhk, edu .k, zhanxi Ahanxiaayyannairr tsinghua. edu. cn
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 arbi- trary new tasks in a zero-shot manner. While the popular Forward-Backwe representations (FB) and related methods have shown promise in zero-shot RI., empirically found that their modeling lacks expressivity and that extrapola- tion errors caused by -uitt--distribution (OOD) actions during offline learning sometimes lead to biased representations, ultimately resulting in suboptimal per- formance. To address these issues, we propose RL with Expressivity enhancement (BREEZE), 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 lusk-conditioned diffusion model, enabling the generation of high-quality and multimodal action distributions in zero-shot RL settings. Moreover, BREEZE employs expressive attention-basee 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: htttttts//gittutt hton//hitterrrr//BEEZZ
REguariizEd Zero-shot
Introduction
Reinforcement learning (RL) has become cornerstone of artificial intelligence, enabling trans- formative advances in robotics (33, autonomous systems [29]. industrial control [71], and large language models (LLM) [60]. However, its real-world adoption faces two persistent challenges: the reliance on human- cronan-provieed reward functions and its task-specific learning paradigm, which novel multiple tasks. These challenges sparked growing interest zero-shot 23
limits adaptability
56 [3], which enables learning a versatile agent through pretraining on reward-free transitions and then zero-shot adaptation arbitrary reward functions during inference. This opens up new possibilities for developing general-purpoot RL agents capable of generalizing across diverse tasks in open-world scenarios.
RL [57] 58
44 26
Existing zero-shot RL approaches mainly fall into two categories:
task/skill-conditioned RL 13
2 and dynamic representation-bant based methods 4 3 demonstrations or reward functions into embeddings as conditioning signals for policy learning.
58
26
56 44). The first category encodes
'Equal contribution. Conresponding author:
Work done during internships at Institute for Al Industry Research (AIR), Tsinghua University.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

This preserves the task generalization capability of policies, but often results in heavy, manually designed pretraining without optimality guaranices. The second category of methods adopts more principled approach by decomposing the problem into dynamic representations that can be recomposed for novel tasks without retraining. Among dynamic representation-basee methods, Forward- wawkwdd representations (FB) 58) have recently attracted notable attention, factorizing occupancy measures into two components: a forward representation that captures policy dynamics and backward representation that encodes global state information. Through offline, unsupervised pretraining, the FB framework learns linearized representations that approximate value functions for arbitrary tasks, holding great promise for zero-shot generalization.
However, despite the elegant theoretical framework provided by the FB representations, our empirical studies reveal that the successor measures learned through existing FB-baser methods are often inconsistent and biased, thereby compromising the stability and overall performance (see Section for details). In this paper, we identify two causes for the shortcomings of existing FB-based methods. First, learning complex approximators and multimodal behaviors requires highly expressive models, which current FB methods lack in both their representations and policy, Second, the offline, unsu- pervised pretraining stage suffers from extrapolation error due to out-of-distribution (OOD) actions, which is a similar problem to offline RL 117 32 (33,,
but exhibits more complex behavior. As our results show, naively integrating value constraints can be ineffective [26]. indicating a need for a more delicate OOD regularization mechanism.
Based on the above observations, we propose Sethavior-REglariiie Zero-shot RL with Expressivity enhancement (BREEZE), novel FB-based framework that simultaneously enhances offline learning stability and zero-shot generalization capability. First, we introduce a behavior-reguularized reformu- lation of FB, which mitigates extrapolation errors while preserving the fidelity of representations. Second, we extract the policy using a task-conditioned diffusion model, enabling high-quality multi- modal action distributions in zero-shot RL settings. Finally, we employ expressive representation networks based on the attention architecture to capture complicated dynamics. We conducted exten- sive experiments on the EXORL benchmark (70) and the D4RL Kitchen dataset 14,, under both full datasets and small-sample data regimes. The results demonstrate that BREEZE achieves the best or near-the sar-the-best performance while exhibiting superior robustness.
2 Preliminary
Reward-free Markov decision process (MDP). Zero-shot RL is typically formulated as reward- free Markov Decision Process (MDP) (51), defined by a tuple M = (S, A, P, 7), consisting of a state space S, action space A, transition kernel P 8 x A A(S), and a discount factor E (0. 1). Given an initial state-action pair (so, (0) E S x A and a policy π S + A(A), we denote Pr( No: 0gg T) probability and [E[-|50; ao; π] an expectation under the state-action trajectories (so, ao .... St, az 1220 generated by sampling St ~ -(.1ss-1.. 1:1) and a (-[ss).. The state transition under
given by P,, (ds'ls) ss) P(ds'[s. a))sa)a(aaass))
Approximate dynamic programming- based RL 47 uses an action-value Q-function, or optionally a state-value V-function. The Q-function and V-function for π starting so, an under given reward function r : S + IR are respectively defined as (5 an) == Ei>o ''E "'Er((s ((((+))) 11+1)|sp Su. au, æ| and V. (So) = : Eio 7 ((((r(s++1) so. x]. The goal of the zero-shot RL problem is to train aask-agnostic policy given an offline dataset D = {st, ags Si+1 \. generated by an unknown behavior policy u(-(s), that can later generalize to any downstream task defined by reward function Teval S + R. i.c. find max, E - ''Teval (5t+1) |so ag. r7. relying only on a small set of reward-labeles samples with no further finetuning.
representations (FB). FB introduces rank-d approximation of the successor measure, which is defined as the expected discounted occurrences of future states 8 E S+ after starting from (So, an) under policy
Md (so, ao, S+)
Σn prrrssu E S+ 80, ag, T) VS+ C S.
(1)
420
Following this definition, Q-function under policy π can be expressed as:
(s,
M (s, a, (s...
(2)
J
2

---

Given representation space RR state distribution p, task vector E R and a policy parameterized by 2, FB expresses M as the product of a forward representation F A R I and backward representation B S + IR, resulting in an optimal policy nz:
M** (8; do, ds+) 22 (so, do, =) (d(s)))
(5()) arg max F( 0 =)_z,
where 2 is defined as = = (r)p(r(s))B( from a few samples with a known reward function r. The corresponding Q-function can be obtained immediately as Q2 a) = F(s, a, =)= 2.
Since the successor measure M satisfies a Bellman- like equation = P + pp M FB
derives the temporal difference (TD) objective (50) on M as:
CFB =
( 84ND
(F(st, ar, 2) B(s+) (5 B(++))²
(4)
-2 B [F(ss, a,, 2) B(s++1)]
and incorporate the following objective on enforce the Bellman property on the Q-function:
CF E (e 1-D (F(s). 2).) (((e+))TEp(BB)) ( (8)+)),
(5)
where B(s++1) BB [ 1-1 is the implicit reward estimation and D denote the data distribution.
Recent extensions of FB incorporate additional regularization improve learning performance. For example, Jeen et al. [26) noted that the use of n. (8t+1) in Eq. (4) during offline learning could introduce OOD extrapolation errors. To resolve this problem, they introduce additional CQL-
style (33] regularizers on M (MCFB) or Q2 (VCFB) in LFB for OOD regularization (i.e., decrease values for OOD actions and increase those for data samples).
Pitfalls of Exiting FB-Based Methods
FB offers elegant theoretical foundation for zero-shot RL; however, empirical findings reveal that FB-based methods often produce inconsistent and biased representations in practice. To illustrate this, evaluate the M and Qa-value distributions derived from the learned F and B representations
the vanilla FB [58]. the more recent MCFB (26), and our method on EXORL benchmark [70] with RND 5 datasets. For each method, we its B representation derive the task vector z under default settings with a fixed number of transitions. The M values computed from episode-start state-action pairs (so, 0o) ~ Po and randomly sampled batch of transitions 8 ~ D, while Q-values are evaluated using policy-induced actions (s). Results across additional environments and with method VCFB 26 are provided in Appendix D.3
Figure M illustrates the value distributions from two tasks in the Walker domain. Although the successor measure R is mathematically positive quantity representing future state occupancy, the representations learned by existing FB-based methods fail to accurately capture this property. notable discrepancies: scale mismatch characterized by enor-
These distributions exhibit mous absolute values, and the presence of a considerable number of invalid negative values. Such inaccuracies in the F and B representations subsequently affect downstream Q-value estimates,
Probabiity Denuity
Probabiiity Dersity
Pobabbiitt Density
8.001
HOU
Probubbiity
8.00
k003
800
0.00
0003:
000
WalkerFip
Stalder-Stard
Walker-Flip
StanndSttad
(a) Distribution of empirically evaluated TT
(b) Distribution of empirically evaluated Q=
Figure Visualization of the empirical R and Q. distributions during evaluation stage. We conduct experiments on FB-based methods on two Walker tasks using EXORL [70] RND [5] dataset, We use the learned F. B representations recover the task vector and compute M F B and Q= F'. z. Both the vanilla FB (58) and MCFB 26 result in proportion of error scaling values.
3

---

creating a consistent prediction bias across the system, which is further worsened by the subop- timal action rollout. The CQL-style regularizer in MCFB partially mitigates the issue, shifting the distributions of M and Q= toward more plausible ranges. However, this correction is in- complete, and significant estimation bias persists, indicating the need for more effective reg- ularization strategy. In comparison, our method generates distributions for = and Q2= -vvaue that more closely match theoretical expectations, and can also result in a high-quality policy. The challenges in estimating = and Q:-
value can be primarily attributed to the limited supervisory signal in the vanilla FB learning
750
750
MLPbased F netwark Freewore n BREEZE
MLP bosed networs
Normalized Return
Normalized Return
network BREEZE
objective (Eq. ( ) which fails to constrain the
scale of the learned representations adequately, Additionally, modeling successor measures and zero-shot policies across all possible task vectors z introduces substantial complexity, necessitat- ing models with high expressive capacity. While the marginal gains from MCFB's 's regularizer suggest its potential utility, more effective for- mulation appears necessary. These observations are further corroborated by the results from our method, which integrates refined OOD regular- ization strategy and employs a more expressive model architecture. As result, the distributions of M and Q= values are primarily confined to reasonable range. Importantly, as illustrated in Figure 2 these improvements are not attributable to increased model capacity alone: simply scaling up the original MLP-baaed FB networks does not yield measurable gains. In contrast, the architectural modifications designed in this work lead to a marked improvement in performance, underscoring the influence of model expressivity on the effectiveness of zero-shot learning.
6:0
600
0.1
#Number of Params (M)
#Number of Params (M)
(a) Impact of F Networks
(h) Impact of B Networks
Figure 2: Investigation of the modeling architecture on the vanilla FB [58]. We utilize the peak observed in the training stage of each run a proxy for the architecture's capability ceiling.
4 Behavior- Zero-Shot RL with Expressivity Enhancement
Motivated by these findings, we propose BREEZE, a framework designed to mitigate the OOD and expressivity issues in FB-based methods. We begin by imposing behavioral regularizer on the offline zero-shot RL problem, followed by the transformation in-sample weighted regression paradigm, which can naturally integrate powerful generative models to enhance policy capability. Recognizing the importance of expressivity for learning arbitrary tasks, we describe our practical architecture for stronger representation modeling.
the learning process into an
4.1 Behavior-regularized Optimization
The main OOD extrapolation issue in the vanilla FB learning losses LFB and LF in Eq. (4) and Eq. 5 lies in the policy-generated actions â,+1 Ty (s((+1)) in ([$(+1, learning. Since is trained to maximize F(s, a, z) 2, it may produce actions that appear optimal under the learned representation but are actually OOD and overestimated, leading to extrapolation errors. Our solution to this is to introduce two designs: 1) stabilize the F( Tz (st-1)), 2) value estimate in the LFB; 2) ensure that T. is better regularized by the dataset samples during policy optimization.
during offline
rehuaior-regularizzed representation guidance.
To stabilize the learning of representation learn- ing, we consider regularized version of Lp in Eq. (5) by introducing the task-conditioned state-value function V with respect to policy म, given by:
Va, (8, =) ==
max aEA s.1. p(a)=)>0
F(s, 0, z)
where z is the task vector and u(aas)) is the behavior policy in the dataset D. Compared with the traditional formulation of the state-value function, this optimization problem aims to learn an optimal task-conditioned V-function solely from dataset samples without explicitly estimating u(als). It can be solved through class of value objectives according to different regularization

---

formulations
67
18
30.
In this work, we leverage the commonly used expectile regression as in
IQL
30
Lv = ((..a)-D,a-22 [ (F(s,a, - V=, (s,))))
where (5(u)= -1-- 0)/u², with 7 >>>>>>7 0.5 serving as the expectile parameter. By minimizing the above equations, we can obtain a well-regularized and nearly optimal state-value function V+, for different tasks. Given the reward function rz, associated with task 2, we have the following modified policy evaluation operator T given by
(T"))) (s, == = Ex [r=((') + yVa, (s),
(8)
This formulation of T implicitly introduces regularization to the Q-function. Following above, we modified Ly in Eq. 55 as following Ly-reg
CF Paa [F((,,,
(9)
By substituting the potentially unstable target Q-approximation F(8++1: ([ (8(++1)), =) with more well-behaved state-value function V=,, we can greatly stabilize the learning of representation, as well as respond to optimality demands simultaneously. In comparison with the direct constraint on unseen value approximation, the behavioral regularization guidance ensures maximum preservation of the representation structure, while being flexible for tuning the degree of conservatism.
Behavior-regularize policy extraction. Having introduced regularization for the value function, now turn to policy learning, which prior work suggests be equally final performance [43]. A good policy should effectively leverage high-quality behavior and further generalize near data distribution. To ensure that the policy is well regularized by the dataset samples, we replace the policy objective Eq. (3) with the following behavior-regularized optimization problem:
more critical for
max B (-1)) [((s, a, ()) (s, s.t. / ma(Sdd = Ys DKL(TzllA) S
(10)
where the term F(s, a, 2), Vz, (8) acts as an advantage function maximization, which is equivalent to maximizing the Q-function. The KL constraint anchors π. the behavior policy je, thereby preventing distributional shift in the offline setting. Given the optimal representation F and Vn,, we can obtain the dlosed-form solution for the constrained optimization objective Eq. (TO) by deriving the Lagrangian objective with respect to the policy and setting its derivative to zero, as shown in Proposition
Daatiiionnhhllll below:
Proposition 1.
The solution the constrained optimization problem in Eq. TO yields an optimal policy of the form (See Appendix for proof):
(s) & p(als) exp (F(s, 2) (a (s,
(11)
where 1/a is the Lagrangian multiplier for the KL constraint.
a acts as a temperature that controls the balance between the regularization strength the behavioral policy and the optimization of the value functions. This closed-form solution elegantly combines the distribution of the behavior policy with the optimal value functions, encouraging the policy to favor high-advantage actions while ensuring all actions remain within the support of the dataset, thus balancing performance with OOD avoidance. With this stable guarantee, we now aim to enhance the policy generalization capability in the next section.
4.2 Policy Extraction via Cask-Conditioned Diffusion Model
To compute the solution of the weighted behavior cloning objective in Eq. 11) in practice, a key challenge arises: how to efficiently model and sample from potentially highly complex and diverse distributions distilled from a re-weighted behavior policy?
Direct weighted behavior cloning with Gaussian policy often fails to capture the complex, multi- modal policy distributions 622
20 which a demand for arbitrary task learning. This motivates the known for their capacity to learn complex distributions through
of diffusion models 48
22
37
5

---

iterative denoising. Rather than relying on guidance- techniques that introduce separate, time-dependent term to steer the sampling process, we draw inspiration from recent advances in weighted regression for diffusion models 7 20 27 72). Building on the theoretical foundation of prior work [72], we formalize Proposition
2 how the optimal policy ny can be extracted using a diffusion model trained with weighted regression objective:
Proposition (Tast-conditioned diffusion policy extraction via weighted regression). The extraction of optimal policy nे in Eq. [1] can be achieved by (i) minimizing the weighted regression loss defined as:
min ( exp (asp(aa(F(s,))) (F(,,,(,, Vn, (((.,)))) c0, 8, -o
(12)
with expectations taken over U(O, ((((,,T), 6~ N(O, 1), and (8,a) D, where ot = a,a + ofe follows the forward process (a (a,a, 1) parameterized by noise schedules 0tz or and Rd denotes the corresponding task; and (ii) sampling from T² by solving the corresponding diffusion ODEs/SDES with the learned ca,z. (See Appendix(A A I the detailed proof.)
With the above objective, we can avoid the need for learning the additional time-dependent guidance term and thereby reduce the complexity, while getting stable policy.
For action selection, we employ a rejection sampling mechanism to boost policy performance. Specifically, we first sample K candidate actions (a(1)... { ~ ==($) through the policy rollout. We then evaluate each candidate using the Q-function (approximated by F(s, a, z) and select the action with the highest value:
max a" arg (E(a()) (+))))) F(s, ((s,a,2)*2
(13)
This two-stage approach ensures the policy balances both conservatism and optimism. The diffusion model generates diverse, in-distribution candidates, while the selection step identifies the action with the highest expected return. This combination of expressive generative modeling and value-based selection is crucial for achieving robust, high peeffrerffrmaece zero-shot generalization.
4.3 Expressivity Enhancement for Representation Modeling
The effectiveness of any particllatilllarlly a highly expressive diffusion policy-in leveraging learned action weights hinges on accurate value estimation. Biased value estimates can degrade policy learning to simple imitation, failing to redistribute probability mass toward superior actions. To fully exploit the representational capacity of our diffusion policy, the underlying value representations must precisely capture complex task and dynamic relationships. We thus introduce enhanced network architectures for the forward (F) and backward (B) representations.
As shown in Section 3 we empirically iden- tify pair of attention- based architectures that improve zero-shot performance over the origi- nal FB implementations. Below, we detail the design of our networks, illustrated in Figure
B
Linear
Lineer
Add Layer norm
Layer norm
FFN
FFN
Layer norm
Add Layer
Forward Network. The forward network encodes state-task and state-action pairs us- ing two separate linear encoders, following Touati et al. Touati et al. [58). The state-task encoder captures aas-cconditioned dynamics, while the state-action encoder extracts agent behavior patterns. These two feature sets cor- respond to semantically distinct concepts in the MDP. To accurately approximate the mea-
Self-ATTN
Multi-eadAAT TTN
V
Q
Q
K
N
Layer normh
Trainformer Blocks
Encoder
Encoder
Figure 3: Architecture of BREEZE
sure M(s, a, s', 2) in Eq. 3 the representation F must integrate information from both behavioral and task contexts, capturing their interdependencies effectively. We therefore model the two encoded features as length-2 embedding sequence and process them through self-attention blocks. This allows bidirectional feature refinement between task conditions and agent behaviors. The resulting representation is projected onto d-dimensional space via linear layers.

---

Table IQM results on ExORL benchmark. report overall performance when tasks perform well. Each value is averaged over 10 rollouts on random seeds.
Dataset
Domain
SF-LAP
FB
VCPB
MCFR
HII.P
BREEZE
Walker
516 97
661 + 10
653 d: 22
659 - 51
665 + 33
693 16
RND
Jaco
18 + 18
32 d: 23
46
41 34
52 :
84 + 14
Quadruped
671
+ 18
674
725
Walker
324 1 24
517 :: 99
487 :: 75
578 +: 35
643 1: 22
637 1: 21
APS
Jaco
39 + 26
22 + 14
20 + 18
22 -
84 + 16
132 16
Quadruped
+:
+ 50
679
698
Walker
382 + 129
650 + 19
611 + 94
602 + 112
715 + 31
663 19
PROTO
Jaco
15 = 14
21 + 26
13 = 12
20 + 21
44 + 19
74 + 26
Quadruped
199
135
216
Walker
239 = 79
338 + 74
268 + 67
268 - 97
461 + 64
463 42
DIAYN
Jaco
32 + 26
22 -
24 +
15 =
52 + 7
78 + 11
Quadruped
207
562
511
= 14
666 +
Table 2: IQM results with 100k-dataset. Experiments on randomly sampled 100k transition data from each dataset the ExORL benchmark
Normalized Return
Dataset
Domain
VCFB
MCFB
BREEZE (ours)
60
Walker
264 33
350
287
525 13
40
RND
7 5
20
Qundruped
176 123
233 52
123 61
474
SF-LAP
Walker
370 +: 66
416 + 10
389 4: 77
539 + 15
Kitchen Partial
APS
Jaco
21 11
14 1 13
29
38 9
80
Normalized Return
Quadruped
340 :::29:9
356
60
Walker
415 d: 19
513 d: 31
463 11
553 1: 18
PROTO
40-
Jaco
16 +
18 4: 12
12 -
29 +: 12
Quadruped
198 111
106 + 103
2410 = 134
181 + 60
20-
Walker
SF-LAP HILP BREEZE
DIAYN
Kitchen Mixed
Jaco
17
20
22 + 15
Figure 4: Normalized scores
Quadruped
295 + 46
288 + 48
286 + 31
416 = 78
on Kitchen environment
Backward Network. As defined in Eq. the backward representation B encodes state-level structural information, serving as global embedding the environment. Intuitively, should employ more complex architecture, maintaining orthogonality and enforcing F alignment. We equip B with stack of standard transformer-bas networks with multi-head attention. The final output is projected onto dl dd-dimensional space, consistent with the forward network.
These architecturally enhanced representations enable BREEZE to model complex relationships more accurately, leading to improved value estimates and policy performance. Further implementation details are provided in Appendix B.4
Experiments
We conduct extensive evaluations of BREEZE against previous offline zero-shot RL algorithms across various challenging domains with distinct tasks. We also present the ablations on component choices and hyperparameters.

---

Environmental setup. Our main experiments are conducted on the ExORL benchmark [70].
which provides variety of datasets collected by several unsupervised RL algorithms 34 We
select datasets collected by 4 algorithms: RND [ APS [38 DIAYN [II], and PROTO 69 The
experiments span 3 domains and tasks (Walker: Stand, Walk, Run, Flip; Jaco: Reach Top/Bottom Left/ight; Quadruped: Stand, Walk, Run, Jump), bringing the total to 48 state-based complex tasks for performance calculation. In addition, we consider four challenging multi-stage tasks in the D4RL [14] Franka Kitchen domain [99 with two datasets (mixed and partial), which require long-horizon sequential manipulation on 4 subtasks. The overall setup involves 2 locomotion domains and manipulation ggoal-reaching) domains. For all condiiond domains, we use the backward representation to calculate the goal features as the task vector 2.
Baselines. We only consider offline algorithms that could perform a zero-shot policy generalization for our evaluation. Our baselines include: 1) SF-LAP [3 successor feature- based method with basic features be Laplacian Eigenfunction (LAP) [64,
vanilla FB [57]. 3) VCFB and MCFB [26]:
FB-based offline algorithms with CQL-style regularizer 3 on successor measure M and Q-value separately for OOD issue avoidance, 4) HILP [44]: state-based representation modeling algorithm keeping iestance-preserving temporal structures in latent space having zero-shot capability.
Evaluation. Our experiments are designed to evaluate two key aspects: zero-shot performance and robustness. To assess the overall zero-shot capability, we experiment on the full transitions of each dataset on ExORL and D4RL Kitchen. For robustness, we first present performance curves across training time to demonstrate learning stability and then conduct experiments on a uniformly subsampled 100,000 transitions to examine performance in a limited and low-quality data coverage scenario, following the evaluation process of Jeen al. [26]. We record the Interquartile Mean (IQM), robust metric that reduces the influence of outliers. We evaluate the checkpoints every 10,000 updates with 10 rollouts for each experiment. On Kitchen, we multiply the returns by 25 for normalization, following Park et al. [4). Further implementation and experimental details are provided in Appendix B
5.1 Experimental Results
We analyze the findings of our experiments through the following key questions:
Does BREEZE possess zero-shot capability for novel tasks? Following previous studies [1,26, we record the highest aggregated mean performance across random seeds in the offline pretrain stage. Table Table[]repote reports this score across all tasks on each domain of ExORL, averaged on 5 random seeds, with standard deviations. BREEZE achieves the best or near-best returns in most domains, demonstrating superior zero-shot generalization capability. Moreover, BREEZE can enhance vanilla FB's generalization capability in long-horizon tasks. As shown in Figure we draw the box with the performance reported in the previous study [44), and record BREEZE's top averaged returns across 4 random sceds. While vanilla FB struggles in these long-horizon tasks, BREEZE achieves significantly higher performance, suggesting its ability to better leverage the dataset for complex sequential manipulation.
More results are provided in Appendix D.2
How does BREEZE quantitatively compare to baselines in terms of stability and convergence speed? We provide the learning curves in RND [5] datasets as Figure 5 In locomotion domains (i.c., Quadruped and Walker), BREEZE converges faster to higher performance, with smoother curves and lower variance (as shown by the small shaded area). In the manipulation domain (i.c., Jaco), BREEZE substantially outperforms all baselines with a higher learning speed. Overall, BREEZE demonstrates faster convergence and enhanced stability across both locomotion and manipulation scenarios.
How does BREEZE facing different quality, diversity of datasets? Table 2 shows results on 100,000- 000.000-transitton subsets of EXORL datasets. BREEZE maintains a pronounced advantage over FB-based baselines in this small-sample regime. Compared explicit constraint methods (MCFB/VCFB), these results highlight the importance of behavior alignment across different data regimes. BREEZE also exhibits superior stability, as shown in Figure 6
Refer to Appendix D.I for more learning curves
8

---

Normalizee Return Normalized Return Normalized Return
quadruped-md
jump
run
stand
walk
600
583
283
204
309
230
Steps
Steps
Steps
Steps
Steps
jaco-md
BL
BR
TL
TR
1.3
1.70
Steps
205
Steps
200
Steps
100
100
Steps
340
walker-rnd
flip
stand
walk
2000
400
380
c20
200
319
430
200
140
Steps
146
mcfb
vcfb
hilp
BREEZE(Ours)
Figure 5: Learning Curves on EXORL RND. The solid lines represent the average return over random seeds, while the shaded area denotes the standard deviation.
Nomalized Return
walker-aps
flip
run
stand
walk
502
430
203
482
200
Steps
Steps
Stops
Stops
Stops
fb
mctb
vctb
BREEZE(OUrS)
Figure 6: Learning Curves on 100k-subsample transitions on EXORL APS Walker domain. The solid lines represent the average return over 3 random seeds, while the shaded area denotes the standard deviation.
Quadruped RND
Quadruped RND
Noraalzzdd Retum
Datuset
w/o FB Enhancement
w/o Diffusion
BREEZE
Normalized Return
Walker-RND
64
707-113
693
Jaco-RND
80=24
62:-7
Quadrped-RND
685-1:3
530::33
Walker-APS
614.148
587.153
637
JaCo-APS
82=13
45:1:50
132
Quadruped-APS
655$220
568=19
Guiddanc Weight
Table Ablation on necessity of FB enhancement and diffusion policy (3 seeds).
Figure 7: Ablation on the temperature.
Figure 8: Ablation on the expectile fac- tor.
5.2 Ablation Studies
Hyperparameter choices. We ablate two key hyperparameters in our behavior-constrainee algo- rithm: the expectile value and the temperature a. As shown in Figure and Figure (Quadruped domain, RND dataset, 3 random seeds), performance improves near monotonically with 7 and peaks = 0.05, and similarly to other domains. This indicates that, while in-sample conservatism crucial, the algorithm benefits from more aggressive, value- weightwe optimization. Consequently, we select 0.99 and a 0.05 as defaults in ExORI., optimally balancing conservatism with performance for downstream tasks.

---

Necessity of each component. We ablate our two core innovations an enhanced representation model and diffusion policy-to assess their necessity RND and APS datasets. BREEZE is ablated using: a gaussian policy variant (w/o Diffiusion) and the baseline architecture variant (w/o FB Enhancement). We process transitions with data normalization for the diffusion policy, as in a commonly used preprocessing step in previous diffusion- based methods. We cancel this process and return to the default policy setting to have a reasonable comparison of the components' contributions. As shown in Table 3 each component somehow individually slightly boosts performance, but their synergy creates a leap in zero-shot capability. This demonstrates a critical interdependence: the design of BREEZE is logical combination to handle the system issue.
Related Works
Unsupervised zero-shot RL. Recent advancements in unsupervised zero-shot RL 34
mainly build upon the utilization of successor representation (SR) [10]. with two main branches: successor feature (SF) 3 and forward-backwart representation (FB) (57 58]. Both approaches require decomposing and linearizing reward-aware dynamics to enable zero-shot inference for arbitrary downstream rewards. Other approaches relax the linearity assumption by modeling SR as full distributions using generative models [25]
55 63 12]. Among all, FB has emerged as a pivotal framework, factorizing successor measures into forward and backward components and avoiding the potential collapse issue in SF at the same time. Extensions like MCFB/ (26) build on offline RL, enforcing explicit value constraints to mitigate the over-estimation issues that arise from updating on OOD rollouts. FB-CPR (56] introduces regularization of the Q-estimation by imitating pre-collected datasets in an online setting. FB-AWARE (6] introduces an autoregressive mechanism to handle the ver-lincarizttion problem in vanilla FB. Unlike SR-based approaches. HILP [44] structures the latent space so that distances correspond to transition times between states, enabling hierarchical, goal-conditioned, and zero-shot policy learning. Mcanwhile, functional-encon methods [13] 23 directly encode task information by pretraining on manually designed reward functions, but are constrained by the need for extensive reward engineering and fine-tuning.
Offiine RL. Numerous studies address offline RL, in which learning is restricted to a fixed dataset without online interaction, leading to distributional shift. The earliest approaches use policy constraints, enforcing the learned policy to stay close to the behavior policy [ 15 355 or the distribution coverage to stay near the dataset [36 g) Another established line of work employs value
regularization, which directly penalizes the value estimates of OOD actions [33 31
68 66).
Instead of an explicit constraint, in-sample learning methods 330 67 65 61 improve stability by learning values and policies only from state-action pairs within the dataset.
Moving beyond conventional Gaussian policies, recent studies have highlighted the importance of modeling multimodal action distributions 662 20 7 Diffusion models [48] 22 37 known for their strong multi-modal modeling capability, have been well studied in imitation and trajectory modeling 24 approaches use the value function as energy guidance for sampling [39 40). Other works, such as SBBC 77 IDQL 220). and EDP p extract diffusion policy via weighted regression under the in-sample learning framework.
73 g [2], and have recently been integrated into offline RL frameworks. Some
7
Conclusion
In this work, we present BREEZE, a novel framework that mitigates the improper scaling issue in existing FB-based zero-shot RL methods. Specifically, BREEZE tackles two critical limitations: offline extrapolation errors and constrained expressivity, Our solution integrates behavior-regularizee value estimation with a lask-conditioned diffusion policy, enabling stable in-sample learning while capturing complex. multimodal action distributions. Coupled with expressive attention- based repre- sentations, BREEZE more accurately models dynamics and value functions. Extensive experiments demonstrate that BREEZE consistently outperforms existing methods across various benchmarks, underscoring the importance of calibrated regularization and sufficient model capacity for zero-shot generalization. While BREEZE demonstrates improved robustness, its primary limitation is the in- creased computational cost from diffussion-aased sampling, which a common trade-off for utilizing high-performance generative policies. Further discussion is provided in Appendix E
10

---

Acknowledgment
This work is supported by the Wuxi Research Institute of Applied Technologies, Tsinghua University, under Grant 20242001 120, and funded by Horizon Robotics, Asialnfo, and the Xiongan Al Institute.
References
[1] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 2021.
[2] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton, Layer normalization, arXiv preprint 167.v01607004400 2016.
[3] André Barreto, Will Dabney, Rémi Munos, Jonathan J. Hunt, Tom Schaul, David Silver, and Hado van Hasselt. Successor features for transfer in reinforcement learning, In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems, 2017.
[4] Diana Borsa, André Barreto, John Quan, Daniel Jaymin Mankowitz, Rémi Munos, H. V. Hasselt, David Silver, and Tom Schaul. Universal successor features approximators. In The Seventh International Conference Learning Representations, 2019.
[5] Yuri Burda, Harrison Edwards, Amos J. Storkey, and Oleg Klimov. Exploration by random network distillation. ArXiv, abs/1810012831 2018.
[6] Edoardo Cetin, Ahmed Touati, and Yann Ollivier. Finer behavioral foundation models
via auto-regressive features and advantage weighting. Reinforcement Learning Conference, abs/2412.04368,, 2025.
[7] Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. In The Eleventh International Conference Learning Representations, 2023.
[8] Peng Cheng, Xianyuan Zhan, Zhihao WU, Wenjia Zhang, Shoucheng Song, Han Wang, Youfang Lin, and Li Jiang. Look beneath the surface: Exploiting fundamental symmetry for sample- efficient offline rl. In Advances in Neural Information Processing Systems, 2023.
[9] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023.
[10] Peter Dayan. Improving generalization for temporal difference learning: The successor repre- sentation. Neural computation, 544)6613-624, 1993.
[11] Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without reward function. In The Seventh International Conference Learning Representations, 2019.
[12] Jesse Farebrother, Matteo Pirotta, Andrea Tirinzoni, Rémi Munos, Alessandro Lazaric, and Ahmed Touati. Temporal difference flows. In International Conference on Machine Learning, 2025.
[13] Kevin Frans, Seohong Park, Picter Abbeel, and Sergey Levine. Unsupervised zero-shot rein- forcement learning via functional reward encodings. In Proceedings of the 41st International Conference on Machine Learning, 2024.
[14] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2021.
[15] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 14:20132-20144 2021.

---

[16] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. International conference on machine learning. PMLR, 2018.
[17] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning
without exploration. In International Conference on Machine Learning, 2018.
[18] Divyansh Garg, Joey Hejna, Matthieu Geist, and Stefano Ermon. Extreme q-learning: Maxent rl without entropy. arXiv preprint arXiv:330..02328, 2023.
[19] Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman, Relay policy learning: Solving long horizon tasks via imitation and reinforcement learning. Conference on Robot Learning (CoRL), 2019.
[20] Philippe Hansen-Fstruuch, llya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit q-learning as an actor-critic method with diffusion policies, 2023.
[21] Longxiang He, Li Shen, Junbo Tan, and Xueqian Wang. Aligniql: Policy alignment in implicit q-learning through constrained optimization, 2024.
[22] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Proceedings of the 34th International Conference on Neural Information Processing Systems, 2020.
[23] Tyler Ingebrand, Amy Zhang, and Ufuk Topcu. Zero-shot reinforcement learning via function encoders. In Proceedings of the 41st International Conference on Machine Learning. PMLR, 2024.
[24] Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In Proceedings of the 39th International Conference Machine Learning, Proceedings of Machine Learning Research.
[25] Michael Janner, Igor Mordatch, and Sergey Levine. Gamma-models: generative temporal difference learning for infinite-horizon prediction. In Proceedings of the 34th International Conference on Neural Information Processing Systems, 2020.
[26] Scott Jeen, Tom Bewley, and Jonathan M. Cullen. Zero-shot reinforcement learning from low quality data. In Advances Neural Information Processing Systems, 2024.
[27] Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. In Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, 2023.
[28] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. In Annual Conference on Neural Information Processing Systems, 2021.
[29] Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A. Al Sallab, Senthil Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey, IEÉE Transactions on Intelligent Transportation Systems, 23(6):4909 4926, 2022. doi: 10. 1 1109/TITS.2021.30..
[30] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning International Conference on Learning Representations, 2021.
[31] Ilya Kostrikov, Jonathan Tompson, Rob Fergus, and Ofrr Nachum. Offline reinforcement learning with fisher divergence critic regularization. In Proceedings the 39th International Conference on Machine Learning, Proceedings of Machine Learning Research, 2021.
[32] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off off-policy q-learning via bootstrapping error reduction. Advances Neural Information Processing Systems, 2019.
[33] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Proceedings of the 34th International Conference on Neural Information Processing Systems, 2020.
12

---

[34] Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang. Lerrel Pinto, and Pieter Abbeel. URLB: unsupervised reinforcement learning benchmark. In Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks 2021.
[35] Jianxiong Li, Xiao Hu, Haoran Xu, Jingjing Liu, Xianyuan Zhan, and Ya-Qin Zhang. Proto: Iterative policy regularized offline-to-online reinforcement learning. 2023.
[36] Jianxiong Li, Xianyuan Zhan, Haoran Xu, Xiangyu Zhu, Jingjing Liu, and Ya-Qin Zhang. When data geometry meets deep function: Generalizing offline reinforcement learning. In International Conference on Learning Representations, 2023.
[37] Sungbin Lim, Eunbi Yoon, Tachyun Byun, Taewon Kang, Seungwoo Kim, Kyungjae Lee, and Sungjoon Choi. Score-based generative modeling through stochastic evolution equations in hilbert spaces. In Proceedings of the 37th International Conference on Neural Information Processing Systems, 2023.
[38] Hao Liu and Abbeel. Aps: Active pretraining with successor features. International Conference Muchine Leurning, 2021.
[39] Cheng Lu, Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li, and Jun Zhu, Contrastive energy prediction for exact energy-guided diffusion sampling in offline reinforcement learning. In Proceedings of the 40th International Conference on Machine Learning. JMLR.org. 2023.
[40] Liyuan Mao, Haoran Xu, Weinan Zhang, Xianyuan Zhan, and Amy Zhang. Diffusion-dice: In- sample diffuusin guidance for offline reinforcement learning. In Advances in Neural Information Processing Systems, 2024.
[41] Fei Jianye Hao, Yao Mu, Yifu Yuan, Yan Zheng, Bin Wang, and Zhixuan Liang. Metadiffuser: Diffusion model conditional planner for offline meta-rl. In International Conference on Machine Learning, Proceedings of Machine Learning Research, 2023.
[42] Haoyi Niu, Shubham Sharma, Yiwen Qiu, Ming Li, Guyue Zhou, Jianming HU, and Xianyuan Zhan. When to trust your simulator: Dynamics-aware hybrid offline-and-onling reinforcement learning. In Advances in Neural Information Processing Systems, 2022.
[43] Seohong Park, Kevin Frans, Sergey Levine, and Aviral Kumar. Is value learning really the main bottleneck in offline rl? Advances in Neural Information Processing Systems, 2024.
[44] Scohong Park. Tobias Kreiman, and Sergey Levine. Foundation policies with hilbert representa- tions. In International Conference on Machine Learning. PMLR, 2024.
[45] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach Vito, Martin Raison, Alykhan Tejani. Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: an imperative style, high-performance deep learning library. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, 2019.
[46] Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the AAAI conference artificial intelligence, 2018
[47] Warren B Powell. Approximate Dynamic Programming: Solving the curses of dimensionality. volume 703. John Wiley & Sons, 2007.
[48] Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamies. In Proceedings of the 32nd International Conference on International Conference on Machine Learning. 2015.
[49] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 1021.
13

---

[50] Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning. (1):9944, 1988.
[51] R.S. Sutton and A.G. Barto. Reinforcement learning: An introduction. IEEE Transactions on Neural Networks, 95):1054-1054, 1998. doi: 10.1 109/TNN. 1998.712192.
[52] Tianyi Tan, Yinan Zheng. Ruiming Liang, Zexu Wang, Kexin Zheng, Jinliang Zheng. Jianxiong
Li, Xianyuan Zhan, and Jingjing Liu. Flow matching-base autonomous driving planning with advanced interactive behavior modeling. arXiv preprini arXiv:2510. /0003 2025.
[53] Chen Tang, Ben Abbatematteo, Jiaheng Hu, Rohan Chandra, Roberto Martín-Martin and Peter Stone. Deep reinforcement learning for robotics: survey of real-world successes. In Proceedings of the AAAI Conference on Artificial Intelligence, 2025.
[54] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq. Timothy Lillicrap, and Martin Riedmiller. Deepmind control suite, 2018.
[55] Shantanu Thakoor, Mark Rowland, Diana Borsa, Will Dabney, Rémi Munos, and André Barreto. Generalised policy improvement with geometric policy composition. In International Conference on Machine Learning, 2022.
[56] Andrea Tirinzoni, Ahmed Touati, Jesse Farebrother, Mateusz Guzek, Anssi Kanervisto, Yingchen Xu, Alessandro Lazaric, and Matteo Pirotta. Zero-shot whole-body humanoid control via behavioral foundation models. In International Conference Representation Learning, 2025.
[57] Ahmed Touati and Yann Ollivier. Learning one representation optimize all rewards. In Proceedings of the 35th International Conference on Neural Information Processing Systems, 2021.
[58] Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does zero-shot reinforcement learning exist? In The Eleventh International Conference on Learning Representations, 2023.
[59] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, 2017.
[60] Shuhe Wang, Shengyu Zhang, Jie Zhang, Runyi Hu, Xiaoya Li, Tianwei Zhang, Jiwei Li, Fei Wu, Guoyin Wang, and Eduard Hovy. Reinforcement learning enhanced 1lms: A survey, 2025.
[61] Xiangsen Wang, Haoran Xu, Yinan Zheng, and Xianyuan Zhan. Offline multi-agent rein- forcement learning with implicit global-to-local value regularization. In Advances in Neural Information Processing Systems, 2023.
[62] Zhendong Wang, Jonathan J. Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023.
[63] Harley Wiltzer, Jesse Farebrother, Arthur Gretton, Yunhao Tang, André Barreto, Will Dabney, Mare G. Bellemare, and Mark Rowland. A distributional analogue to the successor representa- tion. In Proceedings of the 41st International Conference on Machine Learning, 2024.
[64] Yifan Wu, George Tucker, and Ofir Nachum. The laplacian in rl: Learning representations with efficient approximations. In The Seventh International Conference on Learning Representations, 2019.
[65] Haoran Xu, Li Jiang, Jianxiong Li, and Xianyuan Zhan. A policy-guided imitation approach for offline reinforcement learning. In Advances in Neural Information Processing Systems, 2022.
[66] Haoran Xu, Xianyuan Zhan, and Xiangyu Zhu. Constraints penalized q-learning for safe offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2022.
14

---

[67] Haoran Xu, Li Jiang, Jianxiong Li, Zhuoran Yang, Zhaoran Wang, Victor Chan, and Xianyuan Zhan. Offline rl with no ood actions: In-sample learning via implicit value regularization. In International Conference on Learning Representations, 2023.
[68] Rui Yang, Chenjia Bai, Xiaoteng Ma, Zhaoran Wang, Chongjie Zhang, and Lei Han. Rorl: Robust offline reinforcement learning via conservative smoothing. In Advances in neural information processing systems, 2022.
[69] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with prototypical representations. In International Conference on Machine Learning, 2021.
[70] Denis Yarats, David Brandfonbrener, Hao Liu, Michael Laskin, Pieter Abbeel, Alessandro Lazaric, and Lerrel Pinto. Don't change the algorithm, change the data: Exploratory data for offline reinforcement learning, 2022.
[71] Xianyuan Zhan, Haoran Xu, Yue Zhang, Xiangyu Zhu, Honglei Yin, and Yu Zheng. Deepther- mal: Combustion optimization for thermal power generating units using offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2022.
[72] Yinan Zheng. Jianxiong Li, Dongjie Yu, Yujie Yang, Shengbo Eben Li, Xianyuan Zhan, and Jingjing Liu. Safe offline reinforcement learning with tcasibility-guidde diffusion model. In The Twelfth International Conference on Learning Representations, 2024.
[73] Yinan Zheng, Ruiming Liang, Kexin ZHENG, Jinliang Zheng, Liyuan Mao, Jianxiong Li, Weihao Gu, Rui Ai, Shengho Eben Li, Xianyuan Zhan, and Jingjing Liu. Diffusion-based planning for autonomous driving with flexible guidance. In The Thirteenth International Conference on Learning Representations, 2025.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, the main claims made in the abstract and introduction accurately reflect the paper's 's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Yes, the paper discusses its limitations in Appendix including the high computational demand.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
16

---

Answer: [Yes]
Justification: Yes, the paper provides the full set of assumptions and complete proofs for all theorems in the appendix
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Yes, the paper fully discloses all information needed to reproduce the main experimental results. Pseudocode, all experimental settings, hyperparameters, and baseline details are included in the appendix
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
17

---

reproducible, they 5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [Yes]
Justification: Yes, the code will be made available after the open-source approval process is completed.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/nips. wwtpcc/pccppcpcceeeooooooop / seabbbmmssionoo swiniiossiiiiccyyy Policyy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, all hyperparameters and experimental setup details necessary to under- stand the results provided in the appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification: Yes, for all algorithms, several random seeds are used. All results are presented mean + standard deviation.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
18

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, the Appendix C and B provide detailed information about the specific computational devices used for the experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics htttp:://neurige: http://rpp....oppp.....tccccccee
Answer: [Yes]
Justification:
We make sure preserve anonymity.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Yes, the main paper discusses both potential positive societal impacts and limitations of the work.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
19

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Yes, all sources of data and code properly credited in the appendix B with licenses and terms of use clearly indicated.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
20

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer:[NA)
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
21

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
22

---

A Theoretical Proofs
A.1 Proof of Proposition]]
Proposition 1. The closed-form optimal policy for Eq. (T0) follows the form:
(8) X p(a)s) exp (a (F(s, a, Vx, (8, ()),
(14)
where yt denotes the behavior distribution of the dataset.
Proof. Construct the Lagrangian with dual variables A1 and X2 for the equality and KL-divergence constraints, respectively:
[( A1, A2) Bawn, [F(s,a,2) V=. (s,2]] (l (aailalsdde da do ->1 (DkL E).
(15)
Taking the functional derivative of L with respect Tz (als) and setting it zero yields:
ac F(s,a, V=, (s, 2) x2 (Iog (als) log u(als) + +1) A1 = =0.
(16)
Solving above for π. and absorbing constant terms into the normalization factor Z(s,2) gives:
p(als) (s, exp (a (F(s,,,z)* Vo 2)))
(als)
(17)
where a = 1/A2 and the partition function is defined as:
7(8 2) = / p(als)) exp (a (F(s,a,z) = Va, (s, 2))) da,
(18)
This ensures the policy is properly normalized. The proof is completed.
A.2 Proof of Proposition 2
We begin by conducting theoretical review of diffusion models, and then demonstrate how to derive the optimal policy through task-conditioned weighted regression.
Diffussin (probabilistic) models. Diffusion models 48 22 37) are powerful generative models designed to learn and replicate complex distributions. Given N N-sample dataset D [ from an unknown distribution ((no), diffusion models can accurately approximate g ( and generate new samples from the approximation by the following processes:
Forward process. The forward process gradually adds Gaussian noise on sample zo from time 0 to Jr over T timesteps, with transformation distribution follows:
(x, (m)) = N(, layzo, of
(19)
where (2 and o, are predefined noise schedules ensuring @r (xr(w)) zo (xT) as WN 52) for some ō that independent zg-
Backward process. Starting from sr ~ Narr |0, a²T), diffusion models reconstruct the original data Zo by solving diffusion ODE/SDE [49 from T to 0:
dx; [(U), "9 (t)V logg((z)) dt,
(Diffusion ODE)
(20)
(Diffusion SDE) dx, = [f(t) - (2) log qe (xx) dt + (()
(21)
where f(t) d log Os
(1) do?
2 2 ar of
28 and / is standard Wiener process when
di time flows backwards from T to 0.
dt
dt
23

---

Diffusion models train neural network c0(2,, parameterized by 8 with the following objective 22 37:
min E (00.)=0))(0 [[le co 011
(22)
where To qu (20), 2t at dge. The fitted 0 can be use estimate the score function Va, log (x4)) by c0( = - v,, log qe (zz) and substituted into Eq. (2021) for zo generation.
Energy-guide diffusion sampling. To bias sampling toward preferred distributions, the prior energy-guidane framework (39] provides the re-weighted distribution po(an):
po(2o) x (((x)))
(23)
where (2() is the unknown data distribution, (xa) is any form of energy function that encodes human preferences, and f(x)= 2 0 can be any non-negative function. To solve this sampling problem, previous approaches train separate time-dependent classifier & [39|24), which is computationally expensive and may introduce additional training errors. To address this, previous work 72 introduces weighted regression method perform exact energy guidance:
Lemma (Weighted regression as exact energy guidance). We can sample xn ~ po( where Po is the weighted distribution in Eq. (23) by (i) optimizing the weighted regression loss:
min R ((..))))) [f(E(z))) || c0(s,, 1)
(24)
where To ~ qo (o)) I2 = 0 + oee; and (ii) substituting the obtained co into diffusion ODEs/S- DEs (49) solving process.
Task-conditioned diffusion sampling. In this work, we extend the Lemma and provide the following proposition for task-specific sampling:
Proposition (Task-conditioned diffusion policy extraction via weighted regression). The extraction of optimal policy T in Eq. [77 can be achieved by (i) minimizing the weighted regression loss defined as:
min aa 0
exp (a (F(s,,,2) Vv, 2))
E0, -0. 8, 8, 1
and (ii) solving diffusion ODEs/SDEs [49] by substituting ca,z in Eq. 27
Proof. We aim to extract the optimal policy forming as a weighted distribution as follows:
(als) x juaass) exp (a (F(s,a,z)T= V=, (s,
where μ(als) is the behavior distribution in dataset D and a(s, T= Vs, (8,2) 2) is the task- dependent weight function. Set pz (als) as task-conditioned reweighted version of the unknown data distribution qo(als) and substitute the task-dependent weight function into Eq. (23). we have:
p= (als)) (a(ss)) exp (a (F(s, 2)T= (s,))))))
(25)
There exist normalization constant Z(s, 2) f (u(aas) exp (a (F(s, a, a, Vn, (s. 2)) da ensures (als) to be valid distribution. We now aim to sample from the distribution p= (a]s).
Parameterize the denoising process with task-conditioned model 68, s, z,t) and substitute in Eq. (4). we can get a modified task-conditioned weighted regression objective for the diffusion models by replacing the energy function with F(s, a, 2) Vz. (s,
exp (a (F(s,,,2) Vn, (s, 2)) ||c- (l.-) 8, 8,2,,,,,,11 2,1)
min
(26)
According to the Lemma substituting c6 into the diffusion ODEs/SDEs 449 generates data that
sample from pa(s),), thereby recovering n (a| s). The proof is completed.
B Experimental Setups
Our experiments span two benchmarks, EXORI. 70 and D4RI. 14 across two locomotion domains (Walker and Quadruped) and two manipulation goal-reaching domains (Jaco and Kitchen). All experiments we conducted are in state-based tasks.
24

---

(a) Walker
(b) Quadruped
(c) Jaco Arm
(d) Franka Kitchen
Figure 9: Domains. Walker, Quadruped, Jaco, and Franka Kitchen.
B.1 Datasets
ExORL [70]. ExORL consists of datasets collected by several unsupervised RL algorithms 34 on the DeepMind Control Suite [54]. We select datasets collected by four unsupervised RI. algorithms: APS [8), RND [S) PROTO and DIAYN [T] for each the three domains (Walker, Jaco, and
Quadruped). Each domain has four tasks for evaluation at test time.
D4RL 14
Kitchen. D4RL Kitchen datasets consist of robotic arm manipulating trajectories of different tasks in the Franka Kitchen [19 domain. We adopt two datasets: "kitchen Kiichen-partial-v00 partial-vO" and "kicchen-miee-vOO...
Partial: Subtasks involve the microwave, kettle, light switch, and slide cabinet. Tasks are guaran- teed to be solved in subset of the 'partial' dataset.
Mixed: Subtasks involve microwave, kettle, bottom burner, and light switch. No trajectories solve any tasks completely in the 'mixed' dataset.
Dataset preprocessing. BREEZE incorporates standardized data normalization for diffusion mod- els across all training datasets preprocessing step stabilize the learning process and enhance generalization capability.
B.2
Domains
Walker (Locomotion): A bipedal robot with 24-dimensional states (joint positions/velocities)) and 6-dimensional actions. Test tasks include Flip, Run, Stand, and Walk. Rewards combine dense objectives: maintaining torso height (Stand), achieving target velocities (Run/Walk), angular momentum (Flip).
Quadruped (Locomotion): A four-legged robot with 78-dimensional states and 1 12-dimensional actions. Tasks include Jump, Run, Stand, and Walk, with rewards for torso stability and velocity tracking.
Jaco aaaiipulationGGol-rraacht
A 6-DoF robotic arm with 55-dimensional states and 6- dimensional actions. Tasks involve reaching four target positions (Top/Bottom Left/Right) using sparse rewards based on proximity to goals.
Franka Kitchen [19 (Mamipulationn/Gollraaaaaaaii reachingr A robotic manipulation environment control- ling 9-DoF Franka robot with multiple given test-time objectives. The agent must sequentially achieve four subtasks per episode, receiving sparse rewards for each subtask. States include arm joint positions, velocities, torques, and task-specific object features. Test-time goals are defined as proprioceptive states concatenated with target object states.
B.3 Baselines
We utilize the following open-source source codebases for experiments on baselines. We set the batch size to 512 to ensure a fair comparison with other methods, while keeping other hyperparameters at their default values.
For SF-LAP and (57, MCFB and VCFB 26 we utilized the open-source implementation available httph:tt/// /estps:/ggithhe tttb/gtttttttt..e. www/e/eeeeeeeeeeeeeeeeeeeeeeeee zewo- shot-rl
25

---

For HILP [4), we utilized the sounce implementation for base- zero-shot RL setting, which is available at http://ttttbb..oooooooooggggg///
B.4 Architectures
This section details the model architectures used in BREEZE.
F Network. The F network processes inputs (s, and (s, 2) using two separate MLPs, each with two hidden layers of size 1024 and ReLU activation, projecting the inputs into 512-dimensional fea- ture space. The resulting embeddings are concatenated along the sequence dimension and processed through identical processing blocks. Each block consists of self-attention layer (59] followed by a feed-forward network with residual connections, LayerNorm [2]. and dropout (rate=0.1). Specif- ically, the final representation is flattened and passed through two independent output heads, denoted as F'l and F2, each comprising two linear layers with a hidden dimension of 256 that map to the d-dimensional output space.
B Network. The network based on a standard transformer architecture with 8-head multi-head attention [59] and ReLU activation. State embeddings are processed through transformer blocks and projected to d-dimensional space via a linear layer and further scaled to vd by L2 normalization. We observe that sufficiently large B network essential for achieving strong performance on the ExORL benchmark.
Diffusion Policy. The diffusion policy uses a residual-connecttel MLP the noise predictor, following the IDQL [20) implementation. The hidden dimension is set to 1024. The task vector is incorporated via Feature-wise Linear Modulation (FiLM) [46]: two parallel linear layers generate scaling factor y and a shifting factor 3 from z, which modulate intermediate features h as @ h + B.
B.5 Sampling
Following previous studies 58 variable z:
26
we utilize a mixture of two methods for sampling the latent
Sample uniformly within sphere of radius v in Rdd
* Derive by setting Bs)), where ~ D is randomly sampled from the data set.
In our primary experiments, we employ a balanced mix ratio of 50% between the two methods.
B.6 Value Learning
Training details. We implement the V-network with 3-layer MLPs with 512 hidden dimensions and ReLU activation functions. During training, we set the for expectile regression in Eq. 7 shown We use clipped double value learning [16] for both the Q-value and the -aluue
in Section B.7
rejection sampling rather than tuning existing parameters for greedy optimization. Rejection sampling is commonly used in diffusion policies 77 200 21), stable mechanism that selects N actions from the policy to boost with the highest value. Our practical implementation of rejection sampling is in two phases:
Rejection sampling. To improve the zero-shot performance,
Training phase:
To compute the FB loss (Eq. 4. we first sample full transition batch, which contains states, actions, next states, and next actions. hyperparameter Pa. referred as the dataset mixture ratio, determines the proportion of next actions drawn directly from the dataset. The remaining fraction replaced by actions sampled from the policy æz: we generate Kirain candidate (Kuir) via the diffusion policy and select the optimal next action according to: actions 0.......;:0 :0
4 arg max; F( (4.1) 2)
(27)
(t)
Evaluation phase: We sample Keyal candidate actions during evaluation, and select the optimal action according to the Q-value.
Details of practical Pa. Kwain and Keval is provided in Section B.7
26

---

B.7 Hyperparameter
This section provides the detailed hyperparameter setup. In our experiments, the model architecture and basic algorithm hyperparameters remain unchanged, as detailed Table Domain-specific hyperparameters are detailed in Table and Table 6
Table General hyperparameters used for BREEZE
Hyperparameter
Value
Optimizer
AdamW
Representation learning rate
le-4
V network learning rate
3e-4
Policy learning rate
le-4
Discount factor
0.98
Learning steps
1,000,000
Mini-batch
512
General hyperparameters
Representation soft update factor À
0.01
Policy soft update factor X
0.001
Latent dimension d
50
mixing ratio
0.5
inference steps
10,000
Regularization weight for orthogonality loss w
10
Beta schedule
vp
Exponential advantages clip
(-00, 100]
F preprocessor hidden dimension
1024
F preprocessor hidden layers
2
F preprocessor output dimension
512
F preprocessor activation function
relu
F attention blocks
2
F linear layer hidden dimension
256
F dropout
0.1
B nhead
8
B encoder layers
2
B decoder layers
2
Architecture
B d model
256
B dropout
0.1
B feedforward dimension
2048
B activation function
relu
V hidden dimension
512
V hidden layers
2
V activation function
relu
Policy MLP blocks
Policy hidden dimension
1024
Policy activation function
mish
Policy time embedding
learnd
C Pseudo-Code
Algorithm provides the pseudocode. Our implementation uses PyTorch (55) with all experiments conducted on single NVIDIA A6000 GPU. Peak memory usage under 23 GB.
27

---

Table 5: Domain- sppeciici- hyperparameters for BREEZE with full dataset.
Domain-datasct
Kusin
Pa
Koal
expectile
F-reg coef. wy
Diffusion steps T
Temperature
Walker-RND
0.1
64
0.99
0.001
5
0.05
Walker-APS
0.1
0.99
0.001
5
Palker-PROTO
0.1
0.99
0.001
5
0.05
Walker-DIAYN
0.1
64
0.99
0.001
5
0.05
Jaco-RND
64
0.99
0.001
10
0.05
Jaco-APS
0.1
64
0.99
0.0001
5
0.05
Jaco-PROTO
0.001
Jaco-DIAYN
0.1
0.99
0.0001
0.05
Quadruped-RNY
0
64
0.99
0.001
10
0.05
Quadruped-APS
0
64
0.99
0.001
10
0.05
Quadruped-PReR PROTO
0.99
0.001
10
Quadruped-DIAYA
64
0.99
0.001
Kitchen-mixed
16
0.7
0.001
5
0.1
Kitchen-partial
0.2
4
0.7
0.001
5
0.08
Table 6: Domain- specific hyperparameters for BREEZE with small sample dataset.
Domain-dataset
Kuain
Pu
Keal
expectile т
F-reg coef. wy
Diffusion steps T
Temperature a
Walker-RND
0.2
32
0.99
0.0001
10
0.05
Walker-APS
0.2
32
0.99
0.0001
10
0.05
Walker-
0.2
32
0.99
0.001
10
0.05
Walker-DAAYN
0.2
32
0.99
0.001
10
0.05
Jaco-RND>
Jaco-APS
32
0.99
0.0001
10
0.05
Jaco-PROTO
32
0.99
0.0001
10
0.05
Jaco-DIAYN
O
32
0.99
0.0001
10
0.05
Quadruped-RNV
0.2
32
0.99
0.001
Quadruped-AMe
32
0.001
10
0.05
Quadruped-PROTO
0.2
32
0.99
0.001
0.05
Muadruped-DIAYY
0.2
32
0.99
0.0001
10
0.05
Algorithm 1 Behavior-Reguulriized Zero-shot RL (BREEZE)
Require:
Latent dimension d, mini-batch b, gradient steps N, orthogonality coefficient wi, value
regularization coefficient Wig
1: Normalize dataset D; Initialize networks Fø, B V.. diffusion actor π. with network
2: for 1,....N .N do
3: Sample a mini- of transitions {(st, S1+1, 4411)ihem C D
4: Sample mini-batch of (z bejl ~ d
5: Collect {4(+1hejj)) by mixture of data batch and outputs
6:
Representation learning:
7:
Compute LFB 15) and L (1b) using Eq. 4.. Eq. 9
8:
Compute orthogonality regularization loss Contbo
9:
Cortho (v) | EiJejo [BB()s:) B (e,))) (s.)1 ||B (+))
Update F& and B by LFB + wi Lortho + Wn LF
11:
// Value function learning:
Update Vu using Eq. 7
13: // Diffusion policy learning:
Update eo using Eq. 22
15:
end for
D Additional Results
This section provides a comprehensive showcase all experimental results.
Section D.I
provides the learning curve comparisons.
28

---

Section D.2 reports the detailed results on ExORL benchmark.
Section D.3
presents the visualizations of our empirically evaluated value distributions.
D.1 Learning Curves
This section presents the zero-shot learning curves the ExORL benchmark. Results the full dataset and the 100k small-sample dataset are shown in Figure I0 12 and Figure 13 respectively. The solid lines represent the mean IQM returns random seeds, and the shaded regions correspond to the standard deviation.
quadruped-aps
jump
stand
walk
590
Norma lize
666
km
T3
Steps
Steps
Staps
Steps
Normalized remmllized Return Normalzz Return
jump
503
800
230
20e
200
Steps
Steps jump
Steps run
Steps
Ste ps
quadruped-prott
stand
wa lk
300
103
1x0
183
Steps
quadruped-md
jump
stand
walk
000
437
Normalized
100
290
Steps
Steps
Steps
Steps
Steps
mcfb
vcfb
hilp
BREEZE(Ours)
Figure 10: Curves of zero-shot performance on Quadruped domain.
29

---

jaco-aps
BL
BR
TR
201
Normalized lineliizedd Return Norma lized
100
150
310
100.
195-
I00
20
50.
65
0.
Steps
Steps
Steps
jaco-diayn
BL
BR
TR
146
50
30
Steps
Stops
Stops
Steps
Stops
Norma Return
jaco-proto
TR
100
56
Steps
bm
Steps
Stops
1.9
5tn
100
Steps
Normalized Return
jaco-rnd
BL
BR
TL
TR
321
120
100
10
32
Steps
Steps
mcfb
vcfb
hilp
BREEZE(Ours)
Figure 11: Curves of zero-shot performance on Jaco domain.
Norma lized Retur
walker-aps
flip
stand
L048
walk
404
wm
km
200
J50
Steps
Steps
Steps
Steps
Normalzee lized Return
walker-diayn
500
393
1x0
59
200
Steps
Steps
Steps
Steps
Ste ps
Normalizes Return
waakker-proto
flip
run
stand
walk
1068
100
scc-
230
Stops
Lok
Stops
Normalized Return
walker-md
rlip
stand
walk
401
4CC
530
203
10
Steps
Steps
Steps
Steps
Steps
mcfb
vcfb
hilp
BREEZE(Ours)
Figure 12: Curves of zero-shot performance Walker domain.
30

---

quadruped-aps
jum
run
stand
wa lk
Normalized
209
210
203
23
12
05
Steps
Steps
Norma Return Normaliza Return
quadruped-dlayn
jump
stand
lk
200
ccs.
200
Steps
Stops jump
Stops
Steps
Stops
quadruped-prost
run
stand
walk
300
100
Steps
Steps
tce
Stops
Jel
Steps
Stops
Normalized Return
quadruped-rn
jump
run
stand
wa lk
400
300-
200
23
35
0.
02
Steps
Steps
Steps
Steps
fb
mcfb
vcfb
BREEZE(OUrs)
Figure 13: Curves of zero-shot performance on Quadruped domain with 100k dataset.
Norma lized Return
jaco-aps
BL
BR
TL
TR
116
Steps
Steps
Stops
Steps
Normalized Mimalizzdd Return Norma Return
iaco-diayn
130
Steps
Steps BL
Steps
Steps
Steps
iaco-proto
BR
TL
TR
Steps
Stops
Normalized Return
jaco-md
BL
BR
TR
21-
Steps
Steps
Steps
Steps
Steps
mcfb
vcfo
BREEZE(Ours)
Figure 14: Curves of zero-shot performance Jaco domain with 100k dataset.
31

---

Norma Hormalized lieee Return
wa walker-aps
run
stand
wi lk
000
330
aCc-
12
05
0.
Steps
Steps
Normalizer Hermalized Return
wa lellk-ddlayn
flip
stand
lk
200
123
205
100
Steps walker-proto
Stops flip
Stops
Steps
Stops
Norma lize Retur
run
stand
lk
000-
ACC
389
Steps
Steps
Stops
Ice
Steps
Steps
Normalized Return
walker-rnd
flip
run
stand
wa lk
359
ACC-
200
33
Steps
Steps
Steps
Steps
fb
mcfb
vcfb
BREEZE(OUrs)
Figure 15: Curves of zero-shot performance on Walker domain with 100k dataset.
Kitchen-mixed
Kitchen-partial
w
Return
Return
Steps
166
Steps
Figure 16: Curves of zero-shot performance on Kitchen domain.
D.2 Detailed Results on ExORL Benchmarks
This section presents complete experimental results on the ExORL benchmark. Table summarizes performance on the full dataset, while Table [ reports results on the 100k small-sample datasets. Evaluations were conducted every 10,000 steps during offline training. We computed the aggregate score across all tasks as the overall performance measure, and selected the highest domain-level score across 5 random seeds along with corresponding per-task results.
32

---

Table 7: Full dataset experimental results on EXORL
Task
SF-LAP
VCTB
BREEZE
walk
855
918
950
Walker
328
466 107
636
+:
517
whole performenaee
061
650
93
RND
reach top right
2 2::2
76 =: 68
78 1: 63
86 68
43 = 42
108 42
reach top left
2 2
4.1.8
15 1 16
10 1.9
50 17
68 1 30
Juco
reach bottom right
1 0.0
46 1.
81 1 1.62
36 10010.
53 41
77 1. 45
reach bottomm left
70 d: 71
4:2
9 10
11 11:::
60 59
88 =: 65
whole performence
1:
32 00:4:00 20
1: 36
4.1 11:3444
84 1:
64
Quadkgeed
+:
484
jump
whole performannc
330 d: 165
671 14
609 26
684 1: 18
677
725 23
walk
2:
725 118
705 4: 130
817 :: 88
802 22
935
stann
663 #:
720 208
567 + 109
751 = 65
860 10
863
Walker
139
244 58
245
316 # 26
332 26
267
flip
221 95
980 54
431 78
430 97
578 48
482
performence
517
1:75
578
APS
reach top right
1. 22
33
10
1.8
119 5252
reach top left
9 15
16 d:
4:
15 1::4::1
80 =
153 76
Jaco
reach bottom right
+
33
37 1: 46
31
105
70
reach bottom
1: 72
6:1.8
13 1::1:1 111
52 =
185
whole performance
38 1: 26
22 1.
20 10:1111
22 3:
84 16
132 16
walk
158
616 48
601
40
stand
796 186
898
891
9134
12
Quadruped
342 199
471
476 10
440 1.
466 13
472
jump
542 1: 126
686 30
692
690
710
whole performunnc
+: 160
668
walk
352 +: 210
880 = 46
783 + 165
807 109
894 38
923 #:
staed
703 +: 155
936 = 30
K32 86
845 :: 1.55
931 22
858 4:
Walker
FI
202 60
208 =: 27
394 4:
202 =: 69
387 20
228
271 +: 156
477 =b 47
503 4:
4.35 148
647 76
643
whole performenee
382 # 129
650 19
6111 #:
602 112
715 91
663
PROTO
right
12:5:
37 07:::52
top
Juco
right
19
reach
whole performence
15
21 1:
13 12
20 21
44
walk
131 1.32
168 1. 76
136 1.63
235 1:4
118 58
298 1.
stand
293 1:
351 d: 263
239 1.
259 132
329 95
509 55
Quadupee
163 d: 51
125 d: 112
100 45
213 172
175 s7 44
283 03
jump
211 1:
244 :1: 02
265 1: 166
160 101
242 70
407
whole performance
199
107
185
135
216
389 44
wolk
174.1 151
=
329 4 147
604 126
613 :t 103
stand
558 1: 129
146
484 1: 194
420 :: 83
606 = 168
712 06
Walker
1188
+
108
26
152
fiip
272 58
386 76
373
whole performancr
230 2:
336 =: 74
368 1:
268 +: 97
461
463 42
reach right
1.
17 12
L. 16
64 1...
76 1
reach left
18
J.
1:2
55
reach bottom right
+: 42
27
7
1
45
06 1:
reach bottom left
12: 411
28:2:
4:
4::
46
82 1
performance
52.17
walk
196 4: 151
475 d: 77
811 4:47
596 +: 26
497 18
482 11
stund
305 +: 270
766 =: 30
753 4:
865 + 30
944 +:
9944 +
Quadruped
run
130 1. 121
380 1. 22
1.
408 1.
469 2
472
jump
190 1. 158
636 19
576 4.
700 1. 18
769
705 1.25
whole performunce
207 +: 168
562 23
511 87
643 +
670 +:
666 4:
33

---

Table Small sample dataset experimental results on EXORL.
VCFB
MCFB
BREEZE
377 16
536
104
144
224
319 46
236 44
performanco
RND
reach top right
28 1:
10
noach top left
7
Jaco
bettor right
+:
19
reach
whole per formunce
T
1
36 +:
walk
126 d:
164 45
94 d: 08
382 4:
stand
208 1: 183
288 1:73
139 1:
651 1: 37
Quadruped
132 1.
186
89 1. 49
318 1.
jump
178 + 1.50
292 82
148 st
544 +: 37
whole performance
170 4: 123
233 52
123
474 1:
Walker
334
whole per formance
77
+:
APS
reach right
36 +: 29
20 + 24
21
reach top left
3 3.2.2
24
21
Jaco
reach bottom right
1. 60
42 1.
56 1.
08
reach betiom
whole performance
21 1: 17
14 14:1:0 13
29 1: 27
38
waik
200
293
255 + 116
435 4:
stand
470
524 141
456 02
810 1:
Quadnyed
260 10
232 46
222 101
401 +
321 1: 115
356 30
338 1. 179
577 1. 59
whole gerformance
57
318 122
536
+
86
57
Walker
24
whole performance
413 19
513 1:31
463
553 1.
PROTTO
moch top right
7 10
17 +
6:
18
reach top left
50 4:4
14 1:: 12
19 14
51 2:
Jacn
reach bettorn right
0 0::00
22 +
12
27
reach botton
11.00
18 1.
12
whole gerformance
16
18 +
12:7
29
103
327
196 202
185
Quadruped
100
98
206 :: 152
244 173
103 1:
106 1: 103
265 194
291
performnnc
108
240 134
1:
114 16
129
418 119
stand
513 1: 245
1. 256
424 1.
524 1.
Walker
87
79
+
+:
wlole per commancce
202 94
210 81
196 096:4:00 26
330 +:
DEAYN
rosc right
18:2 23
14 +
29 +
5
reach top left
32 35
25 4
9 8
72
Jaco
reach bottom right
15 1: 11
12 1.9
42 :2:1: 58
0.11 20
reach botton left
5 5.1.2
19 1 24
0.1.0
7.1. 26
whole perforaacee
18:2
20 10:::00
22
Quadraped
337
par fommunce
+
34

---

D.3 Empirically Evaluated M and Q2 on Other Tasks
We present empirical visualization results of M and Q2 distributions, shown in Figure 17 19 We trained FB, VCFB, and MCFB on the RND [S] dataset. Specifically, the trained network was utilized for inferring the task latent vector rusk and as the input to both =M and Q2 during the evaluation phase.
Probebilitty Densty
Denelyy
Probebility Densiy
Pobbebiityy Density
0.92
PCLE
8007
Melo
0.37
3.33
0.00
8.003
0,03
0.33
0.000
8.03
0.00
Walker-Flio
WalkerStand
Walker-Fip
Walker-Stand
Probabiityy Density
Probataility Density
Probability Density
Probaabiiiy Dereity
0.00
0.007
2223
0233
0.00
2233
0.00
-00
-00
-008
100
-20
1008
24
-00
Walker-Wak
(a) Distribution of empirically evaluated 5 in Walker
(b) Distribution of empirically evaluated Walker
Figure 17: Visualization of the empirical T and Q= distribution on Walker RND.
Prbabiiity Donsity
Deratty
Probability Density
W0T
Probrbiiity
6.03
0.00
6.000
2391
0.001
Jaco Reach Top Left
20 DE De
Jaco-Reach Left
Jaco-Reah Top Right
Probabiity Dersity
Probaiiiity Dansity.
Probabiiity Dereity
Probabiitty Dereity
Von
8:02
2.35
8.00
6000
0335
0000
6.00
2001
0.01
-00--00--00 20 i8s Reach Left
Joco Reach Bottom Right
-28
Reach Left
200
-200
2s un Reach Bottom Right
(a) Distribution of empirically evaluated T in Jaco
[b) Distribution of empirically evaluated Q₂ in Jaco
Figure 18: Visualization of the empirical = and Q2 distribution on Jaco RND.
0.09
2.099
Probabilty Density
Prebbbiiity Denchy
Ba37
VO'
4/3
Probablity
3,01
6.00
Quadruped-jump
Quadruped-WWa Welk
Quedruped-jump
Quedraped-Walk
0.06
2:89
6.06
0.399
Probability Deraity
Probaabilty Deniity
Dersity
0.00
8.00
0.00
Quadruped-Stan
Quadruped-Run
Quadruped Stand
Quadruped Run
(a) Distribution of empirically evaluated M in Quadruped
(b) Distribution empirically evaluated Q3 Quadruped
Figure 19: Visualization of the empirical " and Q. distribution Quadruped RND.
35

---

E Discussion and Limitation
E1. Hyperparameter Sensitivity and Tuning.
BREEZE introduces additional components compared to vanilla FB, resulting in more hyperparame- ters. While performance is generally robust, provide the following tuning guidance:
Effect of components.
observe The diffusion policy ensures a solid performance lower bound by maintaining stable alignment with any quality dataset.
Representation architectures responsible for high performance ceiling, from T 3
Greedy demand for upstream value learning.
Training cost mainly influenced by the action ratio Pa and the number of rejection sampling candidates Kuain- While the diffusion policy provides conservatism, we find that increasing Pe and wain further improves performance in some domains the cost of additional computation. Near-optimal results can be achieved cost-effectively by reducing the batch size and diffusion timesteps reallocate the computational resources to increase Pa and K
E2. Computation and Performance Trade-off.
limitation of BREEZE its higher computational cost, attributable to the expressive represen- tation networks and diffusion model. We consider this reasonable trade-off given the significant improvements in stability and zero-shot performance. For a comparable and fair comparison, all reported results use batch size of 512, and training for million steps takes approximately 20 hours or less for environments with standard TD updates [50]: configurations with higher Kirin and mixture ratios require up 39 hours. As shown Table 9 BREEZE typically converges within 400k steps, matching or exceeding baselines trained for 1M steps. Future work may explore framework-level optimizations, such replacing diffusion with more efficient flow-based policies to lower costs without compromising performance.
Table Computation cost v.s. performance trade-off. Aggregated scores are averaged over random seeds across datasets. Experiments are conducted on a single NVIDIA A6000 GPU.
Methods
Training Steps (k)
Training Time (h)
IQM Walker
IQM Jaco
IQM Quadruped
Aggregate IQM
200
0.8
480
16
303
266
1000
4.0
542
24
531
366
200
2,4
441
18
361
273
322
1000
12.0
505
26
492
341
200
473
21
347
280
MCFB
400
514
21
469
335
1000
12.0
527
25
551
368
200
4.0-7.8
586
57
550
398
(ours)
400
8.0-15.6
606
75
567
416
36


Reviews:
[{"review_id": "YLQhDrU1Mk", "reviewer": "Reviewer_WnNG", "summary": "BREEZE 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 .", "questions": "- 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.", "limitations": "- 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 .", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "- 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) .", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive comments. Regarding the concerns of the reviewer WnNG, we provide the following responses.\n\n**Question 1 [Compute footprint & efficiency]** \n\n- In our experiments, we train BREEZE on a single NVIDIA A6000 GPU. Peak GPU memory usage varied by tasks, normally below 23,000MB.\n- As BREEZE uses more expressive model architectures, it requires more computational resources for training. Training for 1M steps, as adopted in other baseline methods, takes about 20 hours for most environments. However, as we have reported in Fig. 5, BREEZE can learn much faster, typically reaches good performance within just 100k steps, and enjoys stable convergence. By contrast, other baselines need much larger training steps (typically about 500k steps), and often suffer from unstable convergence (performance drop in later training stages).\n- To demonstrate a more detailed computation cost vs performance trade-off, we report the training time per 100k steps and the corresponding aggregated scores in the following table:\n    | Methods | Training Time (h) per 100k Steps| IQM Walker 100k Steps| IQM Jaco 100k Steps| IQM Quadruped 100k Steps| Aggregate IQM Score 100k Steps |\n    |----------|------------------| ------------------| ------------------| ------------------|------------------|\n    | FB | 0.4 |448±12|17±2|270±15|245±6|\n    | VCFB | 1.2 |373±15|18±2|300±27|231±10|\n    | MCFB | 1.2 |415±16|17±1|270±20|234±8|\n    | HILP | 0.75 |509±9|30±1|394±15|311±6|\n    | BREEZE | 2 |615±7|45±3|542±11|404±4|\n- The scores reported for Walker, Jaco, and Quadruped tasks are averaged over five random seeds and all datasets of this domain. The aggregate score denotes the overall performance by averaging all domains in ExORL. Based on the above table and the learning curve in Fig. 5 of our paper, we can observe that BREEZE reaches much better performance over baselines with a reasonable increase in training cost.\n\n**Question 2 & Limitation 1 [Pixel-based experiments]** \n- We appreciate the reviewer's suggestion. In our paper, we follow other zero-shot RL literature and conduct the experiments on the widely used ExORL/Kitchen benchmarks. We are implementing pixel-based experiments in DMControl and will promptly update the results if we can finish producing results within the discussion period. These results will be included in our final version.\n- Regarding Atari experiments, due to the lack of multi-task settings and relevant datasets, they have not yet been used for zero-shot RL in the task setting.\n\n**Question 3, 4 & Limitation 3 [Futher ablation & analysis]** \n- We appreciate the reviewer's suggestion. We will further conduct relevant experiments and provide analysis on the potential impact of critic bias and diffusion sampling.\n- As for the importance of the diffusion model:\n    - Regarding the Kitchen environment, we have finished ablation on the setting w/o diffusion actor, as shown in the following table:\n  \n    | Setting | Dataset | Average Normalized Return |\n    |---------|---------|-------------------|\n    | with Diffusion | kitchen-partial | 72.5 |\n    | with Diffusion | kitchen-mixed | 75.0 |\n    | w/o. Diffusion | kitchen-partial | 33.3 |\n    | w/o. Diffusion | kitchen-mixed | 37.5 |\n\n    - Combining the above results with Table 2 in our paper, we can clearly observe the benefits of using diffusion policy in BREEZE for all tasks, which helps enhance policy expressivity and enables better capturing complex data distributions.\n\n**Limitation 2 [Parameter sensitivity]** \n- In our ablation experiments shown in Fig. 6, the y-axis labels actually start at 650 and 670 instead of 0. The performance varies in the ranges of [650, 700] for $\\tau\\in$ [0.5, 0.99] and [670, 700] for $\\alpha\\in$ [0.01, 0.5]. The performance variations are actually not as significant as they seem to be.\n- Moreover, if the reviewer checks Tables 3 and 4 in the Appendix of our paper, we used the consistent expectile $\\tau$ (0.99 for all ExORL tasks and 0.7 for Kitchen) and guidance weight $\\alpha$ (0.05 for all tasks) for all our experiments without any hyperparameter tuning. Even using untuned hyperparameters for different tasks, BREEZE still achieves better performance. This also demonstrates the hyperparameter robustness of BREEZE.\n\n**Limitation 4 [Comparison to function-encoder lines]** \n- We appreciate the reviewer's suggestion. We have managed to conduct additional experiments on the baseline FRE [1] for the walker domain in the ExORL benchmark during the short 1-week rebuttal period. We use the official codebase in the FRE paper and modify the evaluation part for rollout in our task environment. The results averaged over three random seeds are shown in the table below, and we will report FRE's scores on other tasks and domains in our final version.\n    | Dataset | Tasks (Walker) | FRE IQM Score| BREEZE IQM Score|\n    |---------|-------|--------------|--------------|\n    | RND | Walk | 187.7±9 | 945±3 |\n    | RND | Run | 88.3±3 | 340±16 |\n    | RND | Stand | 407.3±19 | 950±1 |\n    | RND | Flip | 242.3±5 | 593±18 |\n    | RND | Whole Performance | 232±8 | 703±2 |\n    | APS | Walk | 162.3±9 | 949±4 | \n    | APS | Run | 56.7±11 | 278±9 |\n    | APS | Stand | 376.0±19 | 922±33 |\n    | APS | Flip | 104±5 | 589±15 | \n    | APS | Whole Performance | 175±12 | 685±7 |\n    | PROTO | Walk | 177.3±2 | 923±7 |\n    | PROTO | Run | 115.7±6 | 274±9 |\n    | PROTO | Stand | 341.3±7 | 930±4 | \n    | PROTO | Flip | 187.7±35 | 722±13 |\n    | PROTO | Whole Performance | 190±11 | 711±3 | \n    | DIAYN | Walk | 83.3±12 | 774±21 |\n    | DIAYN | Run | 99.0±7 | 229±10 |\n    | DIAYN | Stand | 336.3±10 | 781±49 | \n    | DIAYN | Flip | 94.0±22 | 395±14 |\n    | DIAYN | Whole Performance | 153.2±2 | 545±12 |\n\n- As shown in the above table, BREEZE consistently outperforms FRE on Walker tasks. Function-encoder ZSRL methods, such as FRE, rely heavily on reward prior assumptions and extensive training/tuning, but they have fewer theoretical guarantees compared to FB-based ZSRL methods. \n\n**References**\n\n[1] Frans, Kevin, et al. Unsupervised Zero-Shot Reinforcement Learning via Functional Reward Encodings. ICML 2024."}, {"review_id": "9SUoqNbcn5", "reviewer": "Reviewer_Ytsj", "summary": "This 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.", "questions": "What 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.", "limitations": "No. 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.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## 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.", "was_revised": true, "final_justification": "The rebuttal effectively addresses the previously raised questions and concerns, and demonstrates careful consideration of the feedback. In light of this, the initially positive assessment is maintained.", "rebuttal": "We thank the reviewer for the constructive comments. Regarding the concerns of the reviewer Ytsj, we provide the following responses.\n\n**Weakness 1 [Clarity & writing]** \nWe thank the reviewer for the suggestion. We will perform thorough language editing and polish the writing in our final version.\n\n**Weakness 2 [Preliminaries & task formalization]**\n- We apologize for the relatively brief introduction on the problem setting of zero-short RL (ZSRL) due to the page limit. In short, ZSRL trains a policy in an unsupervised manner based on a dataset collected from a specified domain without reward. The desired reward function will be specified at test time with the same domain, and the learned agent must immediately produce a good policy without fine-tuning (i.e., zero-shot). We will add more preliminary information in our final version. We also refer the reviewer to the prior papers [1,2,3] in FB-based ZSRL for more detailed information.\n- Regarding the discount factor, it is actually necessary to include the discount factor $\\gamma$ in FB-based ZSRL, as the central element in this problem, the successor measure $M^{\\pi}$ (see Eq.(1) in our paper) is defined as the expected **discounted** occurrences of a future state after staring from a given state-action pair. The FB-based ZSRL methods estimate this value by solving a Bellman-like equation $M^{\\pi}=\\mathcal{P}+\\gamma \\mathcal{P}_{\\pi}M^{\\pi}$ (Eq.(3)), in which $\\gamma < 1$ is also needed to ensure convergence to a fixed point during learning.\n\n\n**Weakness 3 & Question 2 [Ambiguity in zero-shot framing]**\n- We actually share the same feeling on the potential ambiguous term \"zero-short RL\", as it is essentially a class of unsupervised RL (training w/o task/reward specifications) that generalizes to arbitrary tasks at test time. However, the term \"zero-short RL\" has already been used extensively in the existing literature, e.g., [1,2,3,4,5,6], hence we follow this naming convention from prior works in our paper. We thank the reviewer for the suggestion and will include a clearer definition in our final version.\n\n**Weakness 4 & Question 3 & Limitation 1 [Missing limitations & assumptions & discussion on failure cases]**\n\nWe thank the reviewer for this constructive comment. \n- We do have a limitation section in Appendix D.4 in our paper. The main limitation of BREEZE is that it has increased computation demand due to the use of more expressive model architectures. But as it enables noticeable improvements in learning performance and stability, we believe it is a reasonable trade-off.\n- Regarding the assumptions:\n    -  BREEZE inherits all theoretical assumptions from FB, e.g., the MDP assumption, as well as assuming for each task vector $z$, we can find a pair of forward ($F_z$) and backward ($B$) representations such that we can decompose $M^{\\pi_z}$ to be  $F^\\top_zB$.\n    -  The availability and quality of offline data require minimal assumptions, as the goal of ZSRL is to learn the environment structure from arbitrary datasets and handle any downstream tasks. In practice, an offline dataset with higher diversity is likely to provide better ZSRL performance, which has been discussed in the prior paper [2]. \n    -  Regarding the use of KL constraint: The KL-based behavior regularization is widely adopted in offline RL literature to prevent the learned policy from deviating too much from the dataset distribution. We borrowed a similar design in Eq.(9) as the FB-based ZSRL problem suffers from a very similar out-of-distribution (OOD) exploitation problem as in offline RL. A more important benefit of this design is that it can enable a closed-form optimal policy (Eq.(9)), which can be seamlessly integrated with a diffusion policy learning objective (Eq.(10)).\n-  Lastly, we will consider adding additional failure case analysis in the final version of our paper, for example, testing the cases with small-sample and low-cover datasets (the OOD problem will become more severe), as well as testing the possible over-conservative issue due to adding additional behavior regularizations.\n\n**Weakness 5 [Related works]**\n\nWe thank the reviewer's recommendation. Reference [7] provides an effective approach to learn improved representations under partial observability. By contrast, FB-based ZSRL methods focus on learning a special type of successor representation that can facilitate zero-shot policy adaptation. The goals of these two representation learning methods are very different. However, we will add extra discussion with [7] in our final paper.\n\n**Weakness 6 & Limitation 2 [Computational details]**\n\nWe appreciate the reviewer's comments and would like to provide the following details. \n- In our experiments, we train BREEZE on a single NVIDIA A6000 GPU. Peak GPU memory usage varied by tasks, normally below 23,000MB.\n- As BREEZE uses more expressive model architectures, it requires more computational resources for training. Training for 1M steps, as adopted in other baseline methods, takes about 20 hours for most environments. However, as we have reported in Fig. 5, BREEZE can learn much faster, typically reaches good performance within just 100k steps, and enjoys stable convergence. By contrast, other baselines need much larger training steps (typically about 500k steps), and often suffer from unstable convergence (performance drop in later training stages).\n- To demonstrate a more detailed computation cost vs performance trade-off, we report the training time per 100k steps and the corresponding aggregated scores in the following table:\n    | Methods | Training Time (h) per 100k Steps| IQM Walker 100k Steps| IQM Jaco 100k Steps| IQM Quadruped 100k Steps| Aggregate IQM Score 100k Steps |\n    |----------|------------------| ------------------| ------------------| ------------------|------------------|\n    | FB | 0.4 |448±12|17±2|270±15|245±6|\n    | VCFB | 1.2 |373±15|18±2|300±27|231±10|\n    | MCFB | 1.2 |415±16|17±1|270±20|234±8|\n    | HILP | 0.75 |509±9|30±1|394±15|311±6|\n    | BREEZE | 2 |615±7|45±3|542±11|404±4|\n- The scores reported for Walker, Jaco, and Quadruped tasks are averaged over five random seeds and all datasets of this domain. The aggregate score denotes the overall performance by averaging all domains in ExORL. Based on the above table and the learning curve in Fig. 5 of our paper, we can observe that BREEZE reaches much better performance over baselines with a reasonable increase in training cost.\n\n**Question 1 [\"Ground truth\" on Figure 1]**\n- For the successor measure $M^{\\pi}$, as it is an expected discounted probability (see Eq. (1)), it should always be positive. But as we have shown in Fig. 1 as well as more results in Fig. 8-10 in the appendix, both FB and MCFB can produce large negative values for $M^{\\pi}$, which is obviously problematic.\n- As for $Q_z$, by definition of value function in RL, for a reward function with range $[0, r_{max}]$, the theoretical Q-value should fall within the range $[0, r_{max}/(1-\\gamma)]$. In the Walker task in Fig. 1, the actual task reward is in the range [0,1], so we should not expect any negative Q-values. However, similar to the case $M^{\\pi}$, we again find FB and MCFB produce highly problematic Q-values. By contrast, BREEZE produces much more reasonable Q-value estimates with proper scaling.\n\n**Limitation 3 [Domain shift consideration]**\n- The ZSRL problem setting [1,2,4,5,6] only considers zero-shot adapting to different tasks in the same domain (e.g., walk, stand, run, flip tasks for Walker domain), rather than focusing on solving domain shift problems. Future research can explore learning agents that have both task and domain generalization capabilities, but it is not the current focus of our paper.\n\n**Paper Formatting Concerns**\n\nWe sincerely appreciate the reviewer's suggestion. We will revise these issues pointed out by the reviewer in the final version of our paper.\n\n**References**\n\n[1] Ahmed Touati, Jérémy Rapin, Yann Ollivier. Does Zero-Shot Reinforcement Learning Exist? ICLR 2023.\n\n[2] Jeen, Scott, Tom Bewley, Jonathan Cullen. Zero-shot reinforcement learning from low quality data. NeurIPS 2024.\n\n[3] Ahmed Touati, Yann Ollivier. Learning One Representation to Optimize All Rewards. NeurIPS 2021.\n\n[4] Park, Seohong, Tobias Kreiman, Sergey Levine. Foundation policies with Hilbert representations. ICML 2024.\n\n[5] Ingebrand, Tyler, Amy Zhang, and Ufuk Topcu. Zero-Shot Reinforcement Learning via Function Encoders. ArXiv 2401.17173.\n\n[6] Frans, Kevin, et al. Unsupervised zero-shot reinforcement learning via functional reward encodings. PMLR 235, 2024.\n\n[7] Altmann et al., CROP: Towards Distributional-Shift Robust Reinforcement Learning using Compact Reshaped Observation Processing. IJCAL 2023."}, {"review_id": "IICaGoJdHx", "reviewer": "Reviewer_Arsf", "summary": "This 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.", "questions": "1. 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.", "limitations": "The 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.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "The authors have sufficiently addressed my questions during the rebuttal. The paper makes an important contribution towards making zero-shot RL more practical. Therefore, I recommend acceptance.", "rebuttal": "We sincerely appreciate the reviewer Arsf for the positive feedback and the constructive comments on our paper. Regarding the comments of the reviewer, we provide the following responses.\n\n**Weakness 1, 2 [Incremental contribution & extra machineries]**\n- Although our work is built upon the FB framework, we provide the first evidence that conventional FB-based zero-short RL (ZSRL) methods suffer from systematically biased successor measure and Q-value estimations, due to the absence of proper scale supervision on the FB learning objective. We show in our paper that enhancing the model expressivity and properly handling the ODO error exploitation are crucial to solving this problem.\n- It should be noted that the extra machinery introduced in BREEZE is targeted specifically to address the above issues. For example, \n    - We introduce expressive self-attention architecture for $F$ and $B$ to capture the complicated state-action-task relationships and enhance successor measure learning quality.\n    - We introduce an implicit behavior regularization schemes for both $F$ and policy extraction, which enable a closed-form optimal policy (Eq. (9)) and can naturally be modeled using a weighted diffusion regression (Eq.(10)).\n- As the reviewer can notice, the above design is a very logical combination. Removing either component will reintroduce the original limitation in the conventional FB-based method.\n\n**Weakness 1, 2 [Performance comparison with HILP]**\n- The baseline HILP is actually not a FB-based ZSRL method. It is constructed upon a goal-conditioned RL framework, by learning a Hilbert representation for state space, and using the normalized distance between the current state and the goal state under the Hilbert representation as the task vector. The final policy in HILP is learned by running an off-the-shelf offline RL algorithm as a sub-routine.\n- As the underlying mechanism of HILP and FB-based method is very different, thus it is not very suitable to directly compare with BREEZE's variants without FB enhancement or diffusion component. These variants are more suitable for comparison with FB-based methods like FB or VCFB/MCFB, and as we have reported in Table 2, even these incomplete variants could have stronger performance than conventional FB-based methods.\n- Lastly, as we have reported in Table 1 and Fig. 5, in almost all tasks, our complete BREEZE algorithm achieves stronger performance and learning speed as compared to HILP.\n\n**Question 1 [Interpretation of Figure 1]**\n- In Fig. 1, for the successor measure $M^{\\pi}$, as it is an expected discounted probability (see Eq. (1)), it should always be positive. But as we have shown in Fig. 1 as well as more results in Fig. 8-10 in the appendix, both FB and MCFB can produce large negative values for $M^{\\pi}$, which is obviously problematic.\n- As for $Q_z$, by definition of value function in RL, for a reward function with range $[0, r_{max}]$, the theoretical Q-value should fall within the range $[0, r_{max}/(1-\\gamma)]$. In the Walker task in Fig. 1, the actual task reward is in the range $[0,1]$, so we should not expect any negative Q-values. However, similar to the case $M^{\\pi}$, we again find FB and MCFB produce highly problematic Q-values. By contrast, BREEZE produces much more reasonable Q-value estimates with proper scaling.\n\n**Question 2 [Orthonormalization regularization]**\n- We'd like to clarify that orthonormalization loss can only regularize the scale and orthogonal property of $B$, but cannot help to address the scaling problem of the successor measure $M$ and Q-value $Q$. Note that $M$ is modeled as $M=F^\\top B$, solely normalizing $B$ is insufficient to enforce proper scaling of $M$. \n- Also, the FB learning objective Eq.(3) completely depends on the bootstrapped update based on learned $F$ and $B$, without any external scale supervision. This is different from the typical Bellman update used in standard RL, as we have reward values as external scale supervision. Hence, the high quality of $F$ and $B$ estimates becomes essential for unbiased FB learning, which requires sufficient model expressivity and careful handling of the impact of OOD error exploitation.\n\n**Question 3 and 4 [Diffusion guidance and missing citation]**\n- Yes, it is possible to perform guided diffusion by using a weighted term to the noise prediction at each step. We refer the reviewer to the detailed proof of Theorem 2 in a prior work [1]. \n- Paper [1] actually is the Ref[42] in Lemma 1 in Appendix A.2. We apologize for the missing reference. For some unknown reason, this citation was not properly compiled in our submission. We will correct this issue in the final version of our paper.\n\n**References**\n\n[1] Zheng, Yinan, et al. Feasibility-guided safe offline reinforcement learning. ICLR 2023."}, {"review_id": "fOCOxjr3l7", "reviewer": "Reviewer_5D2W", "summary": "This 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.", "questions": "1) 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.", "limitations": "Largely 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?", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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", "was_revised": true, "final_justification": "Some of my questions were answered, e.g. on IQL, tau, regularization, mu(a|s). However, my concerns about why attention should be used here is unresolved, as no motivation is provided other than performance, where there are many other design choices at play. I also find the arguments about the distribution of M to be unconvincing because an architecture could enforce only positive values of M without making other improvements. In any case, there approach still has negative values of M. Lastly, the 5x increase  in training time over FB is very concerning given that this paper relies entirely on empirical evidence. This is an extremely important point in my opinion, since this detail is effectively being swept under the rug since training time does not appear in any plots.", "rebuttal": "We thank the reviewer for the constructive comments. Regarding the concerns of the reviewer 5D2W, we provide the following responses.\n\n**Weakness 1 & Question 5 & Limitation 1 [Comments regarding the use of attention architecture]**\n- We provided a discussion in Section 3 to describe why we need to deprecate standard MLP architectures in FB. This is primarily because learning a successor measure for all possible tasks is actually an extremely complex task, which requires sufficient expressivity for modeling. Moreover, as we have shown in Fig.2, simply scaling up MLP parameters does not provide performance improvement; by contrast, using an attention-based architecture greatly boosts the performance.\n- The reason we choose attention-based architecture in our design is threefold:\n    1. Attention mechanism and the derived transformer architecture have already demonstrated great success in terms of modeling capability and scalability. We want to build our model based on a well-verified architecture.\n    2. As illustrated in Fig. 3, the attention mechanism naturally favors capturing the complex dependencies among different dimensions of state, action, and task vectors.\n    3. In our preliminary exploration, we actually tested other network architectures, such as BRO [1] and OFENet [2]. We found that these structures do not provide meaningful performance improvement over the MLP-based FB architecture. Due to the page limit, we only reported the final architecture that works (i.e., attention-based architecture in Fig. 3), rather than these failed attempts.\n- Regarding the concern from the reviewer that \"the inputs are only 1 or 2 tokens\". Each \"token\" here actually refers to a 512-dim embedding vector encoded by the encoder. It is different from the typical discrete token used in NLP tasks. We apologize for the inaccurate description in our initial submission and have revised it in our final paper.\n\n**Weakness 2 and Question 4 [Regulating towards the exploratory policy & how to get $\\mu(a|s)$]**\n- It should be noted that the original FB objective (Eq. (3)) suffers from similar OOD error exploitation issues as in offline RL, where we have discussed this issue in the last paragraphs in Sections 2 and 3. Without proper OOD regularization, the FB policy learning could suffer from serious bias and learning instability. Some recent FB-based methods like MCFB/VCFB (Ref[7] in our paper) have also recognized this issue and used a much more conservative CQL-based regularization to restrict the learned policy from deviating too much from the data distribution (i.e., the exploratory policy mentioned by the reviewer).\n- By contrast, as described in Section 4.1, we adopt a implicit value/policy regularization scheme in BREEZE, that is, through proper reformulation (e.g., leveraging expectile regression in Eq.(7), and using the closed-form optimal policy in Eq.(9)), **we do not need to explicitly learn the exploratory policy $\\mu(a|s)$ in the dataset or enforce explicit policy constraint**, but instead, we can directly using dataset sample to implicit enforce behavior regularization. This scheme is called in-sample learning in offline RL literature [3,4], which is less conservative as compared to the explicit policy constraint or the CQL-based scheme used in MCFB/VCFB, while also providing better performance and learning stability.\n\n**Weakness 3 & Question 1 [Negative values in FB]**\n- In Section 3, for the successor measure $M^{\\pi}$, as it is an expected discounted probability (see Eq. (1)), it should always be positive. But as we have shown in Fig. 1 as well as more results in Fig. 8-10 in the appendix, both FB and MCFB can produce large negative values for $M^{\\pi}$, which is obviously problematic.\n- As for $Q_z$, by definition of value function in RL, for a reward function with range $[0, r_{max}]$, the theoretical Q-value should fall within the range $[0, r_{max}/(1-\\gamma)]$. In the Walker task in Fig. 1, the actual task reward is in the range $[0,1]$, so we should not expect any negative Q-values. However, similar to the case $M^{\\pi}$, we again find FB and MCFB produce problematic Q-values with large negative values. By contrast, BREEZE produces much more reasonable Q-value estimates with proper scaling.\n- The underlying cause for this phenomenon is that the FB learning objective Eq.(3) completely depends on the bootstrapped update based on learned $F$ and $B$, without any external scale supervision. This is different from the typical Bellman update used in standard RL, as we have reward values as external scale supervision. Hence, the high quality of $F$ and $B$ estimates becomes essential for unbiased FB learning, which requires sufficient model expressivity and careful handling of the impact of OOD error exploitation. These are exactly what BREEZE has addressed. The improvement in the final empirical results also demonstrates the effectiveness of our design.\n\n**Weakness 4 [Minor typo]**\nWe sincerely apologize for the typos and thank the reviewer's correction. We will revise these issues pointed out by the reviewer in the final version of our paper.\n\n**Question 2 [Question regarding Eq.(4) and (5)]**\nThe incorporation of $E[BB^T]$ is introduced in the original FB paper [5] rather than proposed by us, which is also adopted in all later FB-based ZSRL works like [6, 7]. Actually, in FB methods, $B$ is designed as a function that maps reward to a task vector $z$ (i.e., $z=E_s[r(s)B(s)]$), hence $B(s')^\\top E[BB^\\top]^{-1}z$ corresponds to the reward estimate based on task vector $z$. For more details, we refer reviewers to the original FB paper [5, 6].\n\n**Question 3 [Typo in Eq. (6)]**\nWe apologize for this typo. Yes, we only max $V_{\\pi_z}$ over $a\\in A$. We will remove $z\\in Z$ under the max operator in our final version.\n\n**Question 6 [Half-time training on HILP]**\nWe apologize that HILP could only plot half-time curves. We used the official codebase with default environment settings and hyperparameters but encountered an environment error around 600k training steps in every domain. Currently, we are not fully clear about the cause of this problem and are still investigating.\n\nTo ensure the correctness, we rigorously checked HILP's original reported results when plotting curves and recording scores, confirming that our obtained peak performances are consistent with the results reported in the HILP paper. \n\nWe will continue fixing this issue and try to obtain the complete curve of HILP in the final version of our paper.\n\n**Question 7 & Limitation 2 [Concern on performance and computation speed]**\n- We would like to clarify that BREEZE outperforms all baselines across most domains and datasets, as shown in Table 5 in the Appendix. Notably, it achieves strong performance simultaneously on all tasks within each domain. \n- Moreover, if the reviewer checks Fig. 5, BREEZE converges much faster than other baseline methods, typically reaches good performance within just 100k training steps, and enjoys stable convergence. By contrast, other baselines need much larger training steps (typically about 300k-500k steps), and often suffer from unstable convergence.\n- To demonstrate a more detailed computation cost vs performance trade-off, we report the training time per 100k steps and the corresponding aggregated scores in the following table:\n    | Methods | Training Time (h) per 100k Steps| IQM Walker 100k Steps| IQM Jaco 100k Steps| IQM Quadruped 100k Steps| Aggregate IQM Score 100k Steps |\n    |-------|---------| --------| -----------| -----------|---------|\n    | FB | 0.4 |448±12|17±2|270±15|245±6|\n    | VCFB | 1.2 |373±15|18±2|300±27|231±10|\n    | MCFB | 1.2 |415±16|17±1|270±20|234±8|\n    | HILP | 0.75 |509±9|30±1|394±15|311±6|\n    | BREEZE | 2 |615±7|45±3|542±11|404±4|\n- The scores reported for each domain are averaged over five random seeds and all datasets. The aggregate score denotes the overall performance by averaging all domains in ExORL. Based on the above table and the learning curve in Fig. 5, we can observe that BREEZE reaches much better performance with a reasonable increase in training cost.\n\nNote: each experiment was conducted on a single NVIDIA A6000 GPU.\n\n**Question 8 [What happens if $\\tau=1$]**\nWhen $\\tau=1$, expectile regression effectively discards all data, equivalent to learning nothing. In our experiments, we used $\\tau=0.99$ for all ExORL tasks, and $\\tau=0.7$ for all Kitchen tasks (see Section B.7 in our Appendix). \n\n**Limitation 3 [Lacking ablation on IQL loss]**\n- Adding behavior regularization directly on a TD-type loss will require adding explicit policy regularizations, i.e., adding $D(\\pi || \\mu)$ to avoid OOD action exploitation, where $D()$ is some divergence function. As we have explained in the response to Weakness 2, this will require learning an explicit $\\mu$ and can be more conservative. We refer the reviewer to the in-sample learning offline RL method papers [3, 4] for more detailed information regarding the advantage of using IQL-style loss in the offline RL setting.\n\n**References**\n\n[1] Michal Nauman, Mateusz Ostaszewski, Krzysztof Jankowski, Piotr Miłos. Bigger, Regularized, Optimistic: scaling for compute and sample-efficient continuous control. NeurIPS 2024.\n\n[2] Kei Ota, Devesh K. Jha, and Asako Kanezaki. Training Larger Networks for Deep Reinforcement Learning.\n\n[3] Kostrikov, Ilya, Ashvin Nair, Sergey Levine. Offline reinforcement learning with implicit q-learning. ICLR 2022.\n\n[4] Xu, Haoran, et al. Offline RL with no ood actions: In-sample learning via implicit value regularization. ICLR 2023.\n\n[5] Ahmed Touati, Yann Ollivier. Learning One Representation to Optimize All Rewards. NeurIPS 2021.\n\n[6] Ahmed Touati, Jérémy Rapin, Yann Ollivier. Does Zero-Shot Reinforcement Learning Exist? ICLR 2023.\n\n[7] Jeen, Scott, Tom Bewley, Jonathan Cullen. Zero-shot reinforcement learning from low quality data. NeurIPS 2024."}]

Author rebuttal:

Meta-review: The paper proposes BREEZE, an improved Forward-Backward (FB) framework for zero-shot RL that addresses instability, non-expressive representations, and out-of-distribution (OOD) action extrapolation in offline RL. 
Strengths include 1) strong empirical improvements across multiple tasks; 2) well-motivated fixes to FB’s key issues (OOD extrapolation, expressivity, multimodality); 3) clear ablations supporting contributions.
Weaknesses include 1) limited domain diversity (no pixel-based or real-robot benchmarks); 2) some missing theoretical clarity (why attention helps, “zero-shot” definition); 3) weak discussion of limitations and failure cases.
The paper is technically strong, with meaningful improvements to zero-shot RL, but reviewers flagged concerns about compute efficiency, scope of evaluation, and clarity. Overall, reviewers are lean towards accept this paper.
Decision: Accept (poster)