When to ASK: Uncertainty-Gated Language Assistance for Reinforcement Learning

TL;DR

This paper presents ASK, a method that combines smaller language models with reinforcement learning to improve out-of-distribution generalization by selectively querying the language model based on uncertainty assessments.

Contribution

ASK introduces a hybrid approach using uncertainty gating to enhance RL agents' OOD performance without retraining, leveraging smaller LMs efficiently.

Findings

ASK achieves a reward of 0.95 in transfer tasks.

Selective querying maintains efficiency while improving OOD navigation.

Effective hybridization requires careful orchestration, not just simple combination.

Abstract

Reinforcement learning (RL) agents often struggle with out-of-distribution (OOD) scenarios, leading to high uncertainty and random behavior. While language models (LMs) contain valuable world knowledge, larger ones incur high computational costs, hindering real-time use, and exhibit limitations in autonomous planning. We introduce Adaptive Safety through Knowledge (ASK), which combines smaller LMs with trained RL policies to enhance OOD generalization without retraining. ASK employs Monte Carlo Dropout to assess uncertainty and queries the LM for action suggestions only when uncertainty exceeds a set threshold. This selective use preserves the efficiency of existing policies while leveraging the language model's reasoning in uncertain situations. In experiments on the FrozenLake environment, ASK shows no improvement in-domain, but demonstrates robust navigation in transfer tasks, achieving a reward of 0.95. Our findings indicate that effective neuro-symbolic integration requires careful orchestration rather than simple combination, highlighting the need for sufficient model scale and effective hybridization mechanisms for successful OOD generalization.