World models adapt quickly to new tasks with low-rank updates
Amortized Low-Rank Adaptation for Model-Based Reinforcement Learning
Machine LearningArtificial IntelligenceRobotics
Summary
Planning agents use world models to imagine what will happen next based on their actions, but these models can struggle when the environment changes. The authors developed CLAW, a method that helps these models adapt quickly and efficiently to new but related environments using only a little new experience. CLAW does this by training a helper network to generate small adjustments to the original model so it can better predict in new situations without heavy computation. Tests showed it adapts faster and more accurately than previous methods while avoiding common problems like overfitting.
What this means in practice
- •For robotics developers: Enable robots to quickly adjust their models for new tasks or environments with minimal data and computation, improving flexibility in real-world settings.
- •For game ai engineers: Improve game characters that can adapt their behavior to changing virtual environments or player styles efficiently during play.
Authors
Fernando Palafox, David Fridovich-Keil
Abstract
World models let agents plan by predicting the consequences of their actions, but changes in the environment can make them inaccurate. We study the problem of adapting a world model to an unknown test-time environment, drawn from a known environment family, using only a few episodes of interaction. Existing approaches trade off computational cost against expressivity, i.e., the range of models a method can produce. For example, in-context learning is computationally cheap but limited in expressivity, and gradient-based adaptation is expressive but computationally expensive. We present CLAW (Context-conditioned Low-rank Adaptation of World models), which addresses this tradeoff by using a hypernetwork to generate low-rank (LoRA) adapters at test time. During pretraining, we simulate adaptation to a variety of environments and jointly train the hypernetwork and base world model. At test time, we freeze the base model and use a forward pass of the hypernetwork to generate adapters from a small batch of test-time transitions. We evaluate CLAW in locomotion and manipulation environment families that vary in dynamics, embodiment, and reward. We show that, using only seconds of test-time data, CLAW outperforms gradient-based adaptation and in-context learning during online adaptation. We also show that CLAW avoids overfitting in data-scarce regimes, that its advantage comes from the expressive adapters rather than context conditioning, and that pretraining the hypernetwork jointly with the base model outperforms training it post hoc.