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 paper studies how to discover a subset of policies that will become useful for fast adaptation of future tasks. The subspace is defined as the convex hull of some anchor policies which are discovered, and are encouraged to perform well with respect to the extrinsic reward while being as different as possible from each other in the parameter space. The quality of the set is then evaluated by its ability to quickly adapt to new tasks.