Learning a subspace of policies for online adaptation in Reinforcement Learning

1 INTRODUCTION . In recent years , Deep Reinforcement Learning ( RL ) has succeeded at solving complex tasks , from defeating humans in board games ( Silver et al. , 2017 ) to complex control problems ( Peng et al. , 2017 ; Schulman et al. , 2017 ) . It relies on different learning algorithms ( e.g. , A2C - ( Mnih et al. , 2016 ) , PPO - ( Schulman et al. , 2017 ) ) . These methods aim at discovering a policy that maximizes the expected ( discounted ) cumulative reward received by an agent given a particular environment . If existing techniques work quite well in the classical setting , considering that the environment at train time and the environment at test time are similar is unrealistic in many practical applications . As an example , when learning to drive a car , a student learns to drive using a particular car , and under specific weather conditions . But at test time , we expect the driver to be able to generalize to any new car , new roads , and new weather conditions . It is critical to consider the generalization issue where one of the challenges is to learn a policy that generalizes and adapts itself to unseen environments . Different techniques have been proposed in the literature ( Section 6 ) to automatically adapt the learned policy to the test environment . In the very large majority of works , the model has access to multiple training environments ( meta-RL setting ) . Therefore , the training algorithm can identify which variations ( or invariants ) may occur at test time and how to adapt quickly to similar variations . But this setting may still be unrealistic for concrete applications : for instance , it supposes that the student will learn to drive on multiple cars before getting their driving license . In this paper , we address a simpler yet harder to tackle generalization setting in which the learning algorithm is trained over one single environment and has to perform well on test environments ; preventing us from using meta-RL approaches . A natural way to attack this setting is to start by learning a single policy using any RL algorithm , and to fine-tune this training policy at test time , over the test environment ( See red/blue points in Figure 1a ) , but this process may be costly in terms of environment interactions . Very recently , the idea of learning a set of diverse yet effective policies ( Kumar et al. , 2020b ; Osa et al. , 2021 ) has emerged as a way to deal with this adaptation setting . The intuition is that , if instead of learning one single policy , one learns a set of ’ diverse ’ policies , then there is a chance that at least one of these policies will perform well over a new dynamics . The adaptation in that case just consists in selecting the best policy in that set by evaluating each policy over few episodes ( K-shot adaptation ) . But the way this set of policies is built and the notion of diversity proposed in these methods have a few drawbacks : these models increase diversity by using an additional intrinsic reward which encourages the different policies to generate different distributions of states . This objective potentially favors the learning of policies that are sub-optimal at train time . Moreover , these approaches make use of an additional component in the policy architecture ( e.g. , a discriminator ) that may be difficult to tune , particularly considering that , at train time , we do not have access to any test environment and thus can not rely on validation techniques to tune the extra architecture . Inspired by recent research on mode connectivity ( Benton et al. , 2021 ; Kuditipudi et al. , 2019 ) and by ( Wortsman et al. , 2021 ) which aims to learn a subspace of models in the supervised learning setting , we propose to learn a subspace of policies in the parameter space as a solution to the online adaptation in the RL setting ( see Figure 1a ) . Each particular point in this subspace corresponds to specific parameter values , and thus to a particular policy . This subspace is learned by adapting a classical RL algorithm ( PPO and A2C in our case , see Section 3.3 ) such that an infinite continuum of policies is learned , each policy having different parameters . The policies thus capture and process information differently , and react differently to variations of the training environment ( see Figure 1b ) . We validate our approach ( Section 5 ) over a large set of reinforcement learning environments and compare it with other existing approaches . These experiments show that our method is competitive , achieves good results and does not require the use of any additional component of hyper-parameters tuning contrarily to baselines . 2 SETTING . Reinforcement Learning : Let us define a state space S and an action spaceA . In the RL setting , one has access to a training Markov Decision Process ( MDP ) denoted M defined by a transition distribution P ( s′|s , a ) : S × A × S → R+ , an initial state distribution P ( i ) ( s ) : S → R+ and a reward function r ( s , a ) : S ×A → R. A policy is defined as πθ ( a|s ) : S × A → R+ , where θ denotes the parameters of the policy . A trajectory sampled by a policy πθ given a MDPM is denoted τ ∼ πθ ( M ) . The objective of an RL algorithm is to find a policy that maximizes the expected cumulative ( discounted ) reward : θ∗ = argmax θ Eτ∼πθ ( M ) [ R ( τ ) ] ( 1 ) where R ( τ ) is the discounted cumulative reward over trajectory τ . Online adaptation : We consider the setting where the policy trained overM will be used over another MDP ( denoted M̄ ) that shares the same state and action space asM , but with a different dynamics and/or initial state distribution and/or reward function1 . Importantly , M̄ is unknown at train time , and can not be used for model selection , making the tuning of hyper-parameters difficult . Given a trained model , we consider the K-shot adaptation setting where the test phase is decomposed into two stages : a first phase in which the model adapts itself to the new test environment over K episodes , and a second phase in which the adapted model is used to collect the reward . We thus expect the first phase to be as short as possible ( few episodes ) , corresponding to a fast adaptation to the new environment . Let us consider that a model πθ generates a sequence of trajectories τ̄1 , τ̄2 , .... , τ̄+∞ over M̄ , the performance of such a model , is defined as : Perf ( πθ , M̄ , K ) = lim T→∞ 1 T T∑ t=1 R ( τ̄K+t ) ( 2 ) which corresponds to the average performance of the policy πθ over M̄ after K episodes used for adapting the policy . Note that we are interested in methods that adapt quickly to new a test environment and we will consider small values of K in our experiments . In the following , for sake of simplicity , K will refer to the number of policies evaluated during adaptation since each policy may be evaluated over more than a single episode when facing stochastic environments . 3 LEARNING SUBSPACES OF POLICIES . Motivation and Idea : To illustrate our idea , let us consider a toy example where the train environment contains states with correlated and redundant features , in such a way that multiple subsets of state features can be used to compute good actions to execute . Traditional RL algorithms will discover one policy πθ∗ that is optimal w.r.t the environment . This policy will typically use the state features in a particular way to decide the optimal action at each step . If some features become noisy ( at test time ) while , unluckily , πθ∗ particularly relies on these noisy features , the performance of the policy will drastically drop . Now , let us consider that , instead of learning just one optimal policy , we also learn a second optimal policy πθ∗′ , but enforcing θ ∗′ to be different than θ∗ . This second policy may tend to make use of various features to compute actions . We thus obtain two policies instead of one , and we have more chances that at least one of these policies is efficient at test time . Identifying which of these two policies is the best for the test environment ( i.e. , adaptation ) can simply be done by evaluating each policy over few episodes , keeping the best one . Our model is built on top of this intuition , extending this example to an infinite set of policies and to variable environment dynamics . Inspired by Wortsman et al . ( 2021 ) proposing to learn a subspace of models for supervised learning , we study the approach of learning a subspace of policies in the parameter space , and the use of such a model for online adaptation in reinforcement learning . Studying the structure of the parameter space has seen a recent surge of interest through the mode connectivity concept ( Benton et al. , 2021 ; Kuditipudi et al. , 2019 ; Wortsman et al. , 2021 ) and obtain good results in generalization , but it has never been involved in the RL setting . As motivated in the previous paragraph , we expect that , given a variation of the training environment , having access to a subspace of policies that process information differently instead of a single policy will facilitate the adaptation . As a result , our method is very simple , does not need any extra hyper-parameter tuning and achieves good performance . 3.1 SUBSPACES OF POLICIES . Given Θ the space of all possible parameters , a subspace of policies is a subset Θ̄ ⊂ Θ that defines a set of corresponding policies Π̄ = { πθ } θ∈Θ̄ . 1In the experimental study , one training environment is associated to multiple test environments to analyze the ability to adapt to different variations . Since our objective is to learn such a subspace , we have to rely on a parametric definition of such a subspace and consider Θ̄ as a simplex in Θ . Let us define N anchor parameter values θ̄1 , .... θ̄N ∈ Θ . We define the Z-space as the set of possible weighted sum of the anchor parameters : Z = { z = ( z1 , ... zN ) ∈ [ 0 , 1 ] N | ∑ zi = 1 } . The subspace we aim to learn is defined by : Θ̄ = { N∑ k=1 zkθ̄k , ∀z ∈ Z } ( 3 ) In other words , we aim to learn a convex hull of N vertices in Θ . Note that policies in this subspace can be obtained by sampling z ∼ p ( z ) uniformly over Z . The advantages of this approach are : a ) the number of parameters of the model can be controlled by choosing the number N of anchor parameters , b ) since policies are sharing parameters ( instead of learning a set of independent policies ) , we can expect that the learning will be sample efficient . Such a subspace is illustrated in Figure 1a through the `` pentagon '' ( i.e. , N = 5 ) in which angles correspond to the anchor parameters and the surface corresponds to all the policies in the built subspace . K-shot adaptation : Given a subspace of policies Θ̄ , different methods can be achieved to find the best policy over the test environment . For instance , it could be done by optimizing the distribution p ( z ) at test time . In this article , we use the same yet effective K-shot adaptation technique than Kumar et al . ( 2020b ) and Osa et al . ( 2021 ) : we sample K episodes using different policies defined by different values of z that are uniformly spread over Z . In our example , it means that we evaluate policies uniformly distributed within the pentagon to identify a good test policy ( blue star ) . Note that , when the environment is deterministic , only one episode per value of z needs to be executed to find the best policy , which leads to a very fast adaptation .

Summary: The authors propose a method for rapid adaptation from a single training task to an unseen test-task in reinforcement learning. The method optimizes to find a convex subspace, specifically a line, of parameters that minimize the objective in expectation over a uniform distribution. For adaptation, the authors propose to find a convex combination that performs well on the test task.