Test-time Offline Reinforcement Learning on Goal-related Experience

Foundation models compress a large amount of information in a single, large neural network, which can then be queried for individual tasks. There are strong parallels between this widespread framework and offline goal-conditioned reinforcement learning algorithms: a universal value function is trained on a large number of goals, and the policy is evaluated on a single goal in each test episode. Extensive research in foundation models has shown that performance can be substantially improved through test-time training, specializing the model to the current goal. We find similarly that test-time offline reinforcement learning on experience related to the test goal can lead to substantially better policies at modest compute costs. We propose a novel self-supervised data selection criterion, which selects transitions from an offline dataset according to their relevance to the current state and quality with respect to the evaluation goal. We demonstrate across a wide range of high-dimensional loco-navigation and manipulation tasks that fine-tuning a policy on the selected data for a few gradient steps leads to significant performance gains over standard offline pre-training. Our goal-conditioned test-time training (GC-TTT) algorithm applies this routine in a receding-horizon fashion during evaluation, adapting the policy to the current trajectory as it is being rolled out. Finally, we study compute allocation at inference, demonstrating that, at comparable costs, GC-TTT induces performance gains that are not achievable by scaling model size.

Goal-conditioned reinforcement learning can teach an agent to navigate between two arbitrary states (e.g. positions) in an environment. Agents are often trained offline: they only observe pre-collected trajectories, without interacting with the environment. During deployment, the agent is allowed to interact with the environment, but its parameters are traditionally kept fixed. This work tries to specialize agents during deployment by updating their parameters on the fly. The key idea is to observe the current state of the agent and the task it is given, use them to filter the pre-collected trajectories for relevant and high-quality data, and quickly train on this data. Crucially, the filtering step does not require external supervision. This process can be repeated periodically, allowing the agent to specialize to its circumstances.