Partial Models for Building Adaptive Model-Based Reinforcement Learning Agents

TL;DR

This work tackles the challenge that deep model-based RL agents struggle to adapt to local environmental changes. It introduces the concept of partial models, partitioning the state space into regions modeled separately and updated with two implementations (simple and scalable) to maintain accuracy across the whole space while enabling rapid local updates. Across Deep Dyna-Q, PlaNet, and Dreamer, validated on MountainCarLoCA, MiniGridLoCA, ReacherLoCA, and their LoCA1/LoCA2 variants, partial models yield faster and more robust local adaptation, addressing interference, update bias, and quick adaptation. These results advance continual RL by showing how modular, region-specific models can sustain distribution-shift resilience and rapid local learning in complex, perceptual environments.

Abstract

In neuroscience, one of the key behavioral tests for determining whether a subject of study exhibits model-based behavior is to study its adaptiveness to local changes in the environment. In reinforcement learning, however, recent studies have shown that modern model-based agents display poor adaptivity to such changes. The main reason for this is that modern agents are typically designed to improve sample efficiency in single task settings and thus do not take into account the challenges that can arise in other settings. In local adaptation settings, one particularly important challenge is in quickly building and maintaining a sufficiently accurate model after a local change. This is challenging for deep model-based agents as their models and replay buffers are monolithic structures lacking distribution shift handling capabilities. In this study, we show that the conceptually simple idea of partial models can allow deep model-based agents to overcome this challenge and thus allow for building locally adaptive model-based agents. By modeling the different parts of the state space through different models, the agent can not only maintain a model that is accurate across the state space, but it can also quickly adapt it in the presence of a local change in the environment. We demonstrate this by showing that the use of partial models in agents such as deep Dyna-Q, PlaNet and Dreamer can allow for them to effectively adapt to the local changes in their environments.

Figures (14)

• Figure 1: The LoCA setup of wan2022towards. The values in the rewarding regions indicate the reward that is received in the corresponding region. $t_s$ indicates the point in which the phase shift happens.
• Figure 2: An instantiation of the simplest implementation of partial models with three pairs of non-parametric and parametric models. The red arrows indicate the direction of information flow from the non-parametric models to the parametric ones. The parametric model consist of an encoder $e$, network body $nb$, decoder $d$, reward head $rh$ and termination head $th$.
• Figure 3: An instantiation of a scalable implementation of partial models with a non-parametric model consisting of three index lists and a parametric one consisting of three reward heads. The red, blue and green arrows indicate the direction of information flow from the non-parametric model to the parametric one.
• Figure 4: Illustration of the MountainCarLoCA and MiniGridLoCA setups. The solid red lines indicate the T1-zone boundaries in Phase 2 of the LoCA setup.
• Figure 5: Plots showing the learning curves of the deep Dyna-Q agents that are referred to as PM-SimImp, PM-ScImp, REG and LOFO on the (a) MountainCarLoCA and (b) MiniGridLoCA setups. Each learning curve is an average discounted return over 20 runs and the shaded area represents the confidence intervals. The maximum possible return in each phase is represented by a solid black line.