Towards Robust Zero-Shot Reinforcement Learning

AI summary

Overview

Research area: Reinforcement learning — specifically offline, unsupervised zero-shot RL built on Forward-Backward (FB) representations, with applications to locomotion and robotic manipulation.

Technical level: Advanced. The paper assumes familiarity with Markov decision processes, successor measures, temporal-difference learning, offline RL, and diffusion models.

Scope: The paper diagnoses two failure modes of existing FB-based zero-shot RL methods (limited model expressivity and extrapolation error from out-of-distribution actions) and proposes a unified fix — BREEZE — evaluated on the ExORL benchmark and the D4RL Kitchen dataset.

What This Paper Is About

Zero-shot RL aims to pretrain a single, task-agnostic agent from reward-free offline data, so that at inference time it can solve arbitrary new tasks defined by arbitrary reward functions without any further training. The popular Forward-Backward (FB) framework gives this an elegant theory, but in practice its learned successor measures are inconsistent and biased: the authors observe enormous absolute values and large numbers of invalid negative values in the learned representations, which corrupt downstream Q-value estimates. BREEZE ("Behavior-REgularizEd Zero-shot RL with Expressivity enhancement") is the authors' upgraded FB-based framework that attacks both the out-of-distribution (OOD) extrapolation problem and the expressivity problem at once.

Key Contributions

1. An empirical diagnosis of FB's pitfalls. The authors visualize the empirical successor-measure and Q-value distributions of vanilla FB and of MCFB on ExORL RND data, showing a scale mismatch (enormous absolute values) and many invalid negative values, and separately show that simply scaling up the original MLP-based FB networks ("Impact of F Networks" / "Impact of B Networks") yields no measurable gains.

2. Behavior-regularized optimization for FB. They introduce a task-conditioned state-value function V with respect to the policy, fitted by expectile regression (as in IQL), and substitute it for the unstable Q-approximation target inside the FB forward-representation loss. Policy learning is reformulated as a KL-constrained advantage-weighted problem whose closed-form solution is proportional to the behavior policy times an exponential advantage weight (Proposition 1).

3. Task-conditioned diffusion policy extraction. Proposition 2 formalizes how the optimal behavior-regularized policy can be extracted by a diffusion model trained with a weighted regression objective on the denoising error, avoiding a separate time-dependent guidance term. At action selection time, K candidate actions are drawn and the one with the highest F(s,a,z)ᵀz is chosen (rejection sampling).

4. Expressive attention-based representation networks. The forward network encodes state-task and state-action pairs with two separate linear encoders, treats them as a length-2 embedding sequence processed by self-attention blocks, and projects to a d-dimensional space; the backward network uses a stack of transformer blocks with multi-head attention, projected to the same d-dimensional space.

Main Findings

• Biased, ill-scaled FB representations. In experiments on two Walker tasks using ExORL RND data, both vanilla FB and MCFB produce error-scaling values in their empirical M^{π_z} and Q_z distributions. Although the successor measure is mathematically a positive quantity, the learned representations contain considerable numbers of invalid negative values. MCFB's CQL-style regularizer only partially mitigates this; BREEZE's distributions fall more closely within the theoretically expected range.

• Expressivity is not just capacity. Merely scaling up the original MLP-based FB networks produced no measurable gains, whereas the architectural modifications designed in this work led to a marked improvement — indicating the effect comes from model expressivity design, not raw parameter count.

• Strong ExORL results at full-dataset scale (IQM, 5 seeds, 10 rollouts). On RND data, BREEZE scores 693 ± 16 on Walker (vs. vanilla FB 661 ± 10, MCFB 659 ± 51, VCFB 653 ± 22, HILP 665 ± 33, SF-LAP 516 ± 97), 84 ± 14 on Jaco (vs. FB 32 ± 23, MCFB 41 ± 34, VCFB 46 ± 35, HILP 52 ± 21, SF-LAP 18 ± 18), and 725 ± 23 on Quadruped (vs. FB 671 ± 14, MCFB 684 ± 18, VCFB 609 ± 29, HILP 674 ± 28, SF-LAP 330 ± 165). On APS data, BREEZE reaches 132 ± 16 on Jaco (best baseline HILP 84 ± 16) and 698 ± 24 on Quadruped. On PROTO data BREEZE reaches 389 ± 44 on Quadruped and 74 ± 26 on Jaco, and on DIAYN it reaches 463 ± 42 on Walker; HILP is slightly ahead on PROTO Walker (715 ± 31 vs. BREEZE 663 ± 19) and on PROTO Quadruped BREEZE is 389 ± 44 vs. HILP 216 ± 54.

• Robustness in the small-data regime (100k transitions). On 100k-transition subsets, BREEZE leads the FB, VCFB, and MCFB baselines in every reported row: RND Walker 525 ± 13 (vs. FB 264 ± 33, VCFB 350 ± 29, MCFB 287 ± 48), RND Quadruped 474 ± 21 (vs. FB 176 ± 123, VCFB 233 ± 52, MCFB 123 ± 61), APS Quadruped 556 ± 52, PROTO Walker 553 ± 18, DIAYN Quadruped 446 ± 78. The authors attribute this to the importance of behavior alignment across different data regimes, contrasting it with explicit-constraint methods (MCFB/VCFB).

• Better stability and convergence speed. On ExORL RND datasets, the learning curves show BREEZE converging faster to higher performance with smoother curves and lower variance in the locomotion domains (Quadruped and Walker), and substantially outperforming all baselines with a higher learning speed in the manipulation domain (Jaco).

• Long-horizon Kitchen improvement. On four multi-stage D4RL Kitchen tasks with the mixed and partial datasets, BREEZE achieves significantly higher performance than vanilla FB, which the authors say struggles in these long-horizon tasks. Returns are multiplied by 25 for normalization following Park et al.

• Component synergy. Ablations (3 seeds) show each component provides only a modest individual boost, but together they create a large leap: on Walker-RND, "w/o FB Enhancement" gives 646 ± 18 and "w/o Diffusion" gives 707 ± 13 versus BREEZE's 693; on Quadruped-RND, 685 ± 13 and 530 ± 33 versus 725; on Jaco-APS, 82 ± 13 and 45 ± 50 versus 132.

• Hyperparameters matter. Performance improves near-monotonically with the expectile τ and peaks at temperature α = 0.05 (on Quadruped, RND dataset, 3 seeds). The authors pick τ = 0.99 and α = 0.05 as defaults for ExORL, balancing in-sample conservatism against value-weighted optimization.

Methodology in Plain English

The starting point is the FB framework, which writes a policy's successor measure (roughly, how often and how soon future states are visited) as the inner product of a forward representation F(s,a,z) and a backward representation B(s′), where z is a task vector derived from a reward function. Training uses a Bellman-style temporal-difference loss.

The authors' first fix targets a specific term in that loss: F(s_{t+1}, π_z(s_{t+1}), z)ᵀz, which is evaluated at actions the policy itself proposes. Offline, the policy may propose actions that look good under the learned representation but are actually unseen in the data — the classic extrapolation-error problem. Instead of clamping values with explicit constraints, the authors learn a task-conditioned state-value function V using expectile regression restricted to dataset actions, and use it as the bootstrapping target. This keeps learning in-sample while still pushing toward optimality.

Second, they reformulate policy learning as maximizing the advantage subject to a KL constraint that keeps the policy anchored to the behavior policy in the data. The closed-form optimum is the behavior policy reweighted by the exponential of the advantage. Because such reweighted distributions can be multimodal and complex — and Gaussian policies often fail to capture that — they fit a task-conditioned diffusion model by weighted regression on the denoising objective, weighting each sample by that same exponential advantage term. At evaluation, they sample several candidate actions from the diffusion policy and pick the one with the highest estimated value.

Third, because a diffusion policy can only be as good as the values steering it, they replace the original MLP representation networks with attention-based ones: self-attention over the state-task and state-action embeddings for the forward network, and a transformer stack for the backward network.

Evaluation covers 3 domains and 12 tasks on ExORL (Walker: Stand, Walk, Run, Flip; Jaco: Reach Top/Bottom Left/Right; Quadruped: Stand, Walk, Run, Jump), giving 48 state-based tasks, using datasets collected by RND, APS, DIAYN, and PROTO. Baselines are SF-LAP, vanilla FB, VCFB, MCFB, and HILP. Checkpoints are evaluated every 10,000 updates with 10 rollouts, and results are reported as Interquartile Mean (IQM) averaged over 5 random seeds (3 seeds for some ablations). Robustness is tested both through learning curves across training time and through uniformly subsampled 100,000-transition datasets.

Why This Matters

Impact on research. The paper argues that the bottlenecks in zero-shot RL are not theoretical but practical — representation bias and OOD extrapolation — and shows that a behavior-regularized, in-sample reformulation combined with expressive generative policies addresses both. It also provides direct counter-evidence to the assumption that more parameters alone will fix FB, which redirects effort toward architecture and regularization design.

Real-world applications:

• General-purpose robot control: pretrain once from reward-free interaction logs, then deploy the same agent on new locomotion or manipulation goals without retraining.
• Goal-reaching and navigation in the Jaco-style domains, where a task vector is derived from the backward representation rather than a hand-specified reward.
• Multi-stage manipulation, such as the long-horizon Kitchen subtask sequences, where sequential decision-making is required.
• Offline industrial control settings where collecting reward-labeled data is expensive but large unlabeled transition datasets exist.

Industry relevance. The setting matches industrial reality: large datasets of unlabeled transitions are cheap, while reward engineering and retraining per task are expensive. A framework that extracts many tasks zero-shot from one pretrained model, and that remains stable when data is subsampled to 100k transitions, is directly relevant to robotics and control pipelines. The authors include an affiliation with Huawei Noah's Ark Lab, and an official implementation is released at https://github.com/Whiterrrrr/BREEZE.

Future Directions

• Extending the behavior-regularized formulation beyond the offline setting to online or hybrid data collection, where the behavior policy changes over time.
• Testing the attention-based representation and diffusion extraction on domains beyond the ExORL locomotion/goal-reaching tasks and the four Kitchen tasks reported here.
• Determining whether the components remain complementary at larger model and data scales, given that simply scaling MLP-based FB networks showed no measurable gain.
• Clarifying the sensitivity of the method to its two key hyperparameters — the expectile τ and the temperature α, which peaked at α = 0.05 — across dataset qualities and domains not studied in the ablations.

Target Audience

Researchers and graduate students working on offline reinforcement learning, zero-shot or unsupervised RL, and successor-representation methods; practitioners building pretrained generalist control policies from large unlabeled datasets; and readers specifically interested in combining diffusion policies with value-based representation learning. A working knowledge of FB representations, expectile regression/offline RL (CQL, IQL) and diffusion models is needed to follow the derivations.

Authors’ 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.