Zero-Shot Reinforcement Learning from Low Quality Data

NeurIPS 2024

Scott Jeen\(^{\alpha}\), Tom Bewley\(^{\beta}\) & Jonathan M. Cullen\(^{\alpha}\)

\(^{\alpha}\) University of Cambridge \(^{\beta}\) University of Bristol

Motivation

• Training policies to (zero-shot) generalise to unseen tasks in an environment is hard! [1]

• Behaviour Foundation Models (BFMs) based on forward-backward representations (FB) [2] and universal successor features (USF) [3], provide principled mechanisms for performing zero-shot task generalisation

• However, BFMs assumed access to idealised (large & diverse) pre-training datasets that we can’t expect for real problems

• Can we pre-train BFMs on realistic (small & narrow) datasets?

[1] Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. JAIR 2023

[2] Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does zero-shot reinforcement learning exist? ICLR 2023

[3] Seohong Park, Tobias Kreiman, and Sergey Levine. Foundation Policies with Hilbert Representations. ICML 2024

Out-of-distribution Value Overestimation in BFMs

Out-of-distribution Value Overestimation in BFMs

Out-of-distribution Value Overestimation in BFMs

[3] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. NeurIPS 2019

Conservative BFMs

Conservative BFMs

Conservative BFMs

Conservative BFMs

Conservative BFMs

[4] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. NeurIPS 2020

ExORL Results

Baselines

• Zero-shot RL: FB, SF-LAP [5]
• Goal-conditioned RL: GC-IQL [6]
• Offline RL: CQL [7]

Datasets

[5] Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does zero-shot reinforcement learning exist? ICLR 2023

[6] Seohong Park, Dibya Ghosh, Benjamin Eysenbach, and Sergey Levine. Hiql: Offline goalconditioned rl with latent states as actions. NeurIPS 2023.

[7] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. NeurIPS 2020

ExORL Results

D4RL Results

Performance on Idealised Datasets is Unaffected

Conclusions

• Like standard offline RL methods, BFMs suffer from the distribution shift

• As a resolution, we introduce Conservative BFMs

• Conservative BFMs considerably outperform standard BFMs on low-quality datasets

• Conservative BFMs do not compromise performance on idealised datasets

Project page:

Twitter/X: @enjeeneer

Website: https://enjeeneer.io