Towards Robust Zero-Shot Reinforcement Learning

TL;DR

The paper tackles the brittleness of zero-shot RL methods that rely on Forward-Backward representations by identifying extrapolation and expressivity issues in offline learning. It proposes BREEZE, a behavior-regularized, expressivity-enhanced FB framework that reframes learning as stable in-sample regression, uses a task-conditioned diffusion model to extract multimodal policies, and leverages attention-based F and B networks for richer dynamics modeling. The key contributions are (i) a novel behavior-regularized objective with expectile-based value guidance, (ii) a diffusion-based, task-conditioned policy extractor with weighted regression and rejection sampling, and (iii) expressive attention-based forward and backward representations, all validated on ExORL and D4RL Kitchen with superior robustness and competitive zero-shot performance. The findings indicate that calibrated regularization and sufficient model capacity significantly improve zero-shot generalization, at the cost of higher computation, suggesting a favorable trade-off for practical open-world RL applications.

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.

Figures (19)

• Figure 1: Visualization of the empirical $M^{\pi_z}$ and $Q_z$ distributions during evaluation stage. We conduct experiments on FB-based methods on two Walker tasks using ExORL yarats2022dontchangealgorithmchange RND Burda2018ExplorationBR dataset. We use the learned $F$, $B$ representations to recover the task vector $z$ and compute $M^{\pi_z}=F^\top B$ and $Q_z=F^\top z$. Both the vanilla FBTouati2022DoesZR and MCFBJeen2023ZeroShotRL result in a proportion of error scaling values.
• Figure 2: Investigation of the modeling architecture on the vanilla FBTouati2022DoesZR. We utilize the peak observed in the training stage of each run as a proxy for the architecture's capability ceiling.
• Figure 3: Architecture of BREEZE
• Figure 4: Normalized scores on Kitchen environment
• Figure 5: Learning Curves on ExORL RND. The solid lines represent the average return over 5 random seeds, while the shaded area denotes the standard deviation.

Theorems & Definitions (7)

• Proposition 1
• Proposition 2: Task-conditioned diffusion policy extraction via weighted regression
• Proposition 1
• proof
• Lemma 1: Weighted regression as exact energy guidance
• Proposition 2: Task-conditioned diffusion policy extraction via weighted regression
• proof