PAnDR: Fast Adaptation to New Environments from Offline Experiences via Decoupling Policy and Environment Representations

1 Introduction

Reinforcement learning (RL) has achieved great successes in many fields, e.g., Game Playing [Badia et al., 2020], Software Testing [Zheng et al., 2019; Zheng et al., 2021], Robotics [Andrychowicz et al., 2020], Industrial Supply Chain [Ma et al., 2021] and so on. Despite of the successes, conventional RL is well known to be lack of generalization ability [Packer et al., 2018]. To be specific, a policy well optimized in a specific environment can be catastrophically poor even in similar environments. The poor generalization among environments and policies severely prevents more practical applications of RL algorithms.

∗Equal contribution.

† Jianye Hao and Yan Zheng are the corresponding authors.

In the literature on RL generalization [Kirk et al., 2021], many approaches like data augmentation [Yarats et al., 2021] and domain randomization [Peng et al., 2018] are proposed to increase the similarity between training environments and the testing environments (i.e., the ones to generalize or adapt to). These methods often require knowing the variations of environment and having the access to environment generation. In contrast, some other works aim to achieve fast adaptation in testing environments without the need for such access and knowledge. Gradient-based Meta RL, e.g., MAML [Finn et al., 2017], aims to learn a meta policy which can adapt to a testing environment within few policy gradient steps. Also in this branch, context-based Meta RL, e.g., PEARL [Rakelly et al., 2019], leverages a context-conditioned policy which allows to achieve adaptation by the generalization among contexts. Useful context representation is learned in training environments to capture the variations; then the context of a testing environment can be inferred from few probed interaction experiences [Fu et al., 2021].

Most prior works consider an online training setting where the agent is allowed to interact with training environments arbitrarily within the interaction budget. However, in realworld problems, online interaction is usually expensive while offline experiences are often available and relatively sufficient. For example, in E-commerce, adapting the advertising strategies obtained from massive offline ads-users interaction data to a new user can effectively save the budgets of advertisers. On this point, a recent work [Raileanu et al., 2020] studies fast adaptation under an offline-training-onlineadaptation setting, where only offline experiences collected on training environments are available before online policy adaptation in new environments. This is a more practical and general setting which is meanwhile more challenging, since the policy learning and variation capture need to be conducted in an offline manner. To solve this problem, Raileanu et al. [2020] propose an extended value function, called PolicyDynamic Value Network (PDVF). By definition, a PDVF additionally takes an environment representation and an policy representation as inputs, and evaluates the corresponding values. In principle, the explicit representation allows the generalization of value estimation among environments and policies, which can be utilized to realize policy adaptation. To be specific, given the offline experiences generated from different combinations of training policies and training environments, the representations of environment and policy are learned from transitions and state-action sequences by dynamics prediction and policy recovery respectively. Thereafter, a quadratic form of PDVF is approximated, according to which the policy is optimized during online adaptation with the environment representation inferred from few interactions in the testing environment. The first drawback is that the representations of environment and policy can have redundant and entangled information since they are from the same interaction experiences of the corresponding combination. It negatively affects the generalization conveyed by the representations which is vital to fast adaptation. Another limitation is that the quadratic approximation of PDVF inevitably makes a sacrifice in expressivity for convenient closed-form optimization, in turn this cripples the optimality of policy obtained during adaptation.

In this paper, we focus on the offline-training-onlineadaptation setting and follow the paradigm of [Raileanu et al., 2020]. We propose Policy Adaptation with Decoupled Representations (PAnDR) for more effective fast adaptation of RL. The conceptual illustration of PAnDR is shown in Fig. 1. Compared to transition-level dynamics prediction adopted by [Raileanu et al., 2020], we adopt trajectory-level context contrast for environment representation learning. By this means, the variation is captured in an integral manner and the inter-environment difference is emphasized, thus offering more useful information and better robustness. One can link our comparison here to that between reconstruction and contrastive learning, where the superiority of the latter has been widely demonstrated [He et al., 2020; Laskin et al., 2020; Fu et al., 2021]. Moreover, we propose a informationtheoretic representation learning method to refine environment and policy representations for information decoupleness and completeness. This is achieved by mutual information minimization and maximization among environment and policy representations and their joint. In addition, distinct to the quadratic form in [Raileanu et al., 2020], PAnDR leverages a typical multi-layer nonlinear neural network for the approximation of the PDVF. Then during online adaptation, the policy is optimized by gradient ascent. In this sense, PAnDR gets rid of such constraints to improve value approximation and policy adaptation. In experiments, we evaluate PAnDR in a series of fast adaptation tasks and analyze the contribution of each part of proposed methods.

We summarize our main contributions as follows: 1) We leverage trajectory-level contrastive learning and propose a novel mutual information based approach for better environment and policy representations. 2) We get rid of the constraint of prior quadratic approximation of PDVF and demonstrate the effectiveness of nonlinear approximation in improving online policy adaptation. 3) We show that PAnDR consistently outperforms representative existing algorithms in our experiments, sometimes with a large margin.

2 Preliminaries

First, we briefly introduce the necessary preliminaries. A detailed discussion on related works is provided in Appendix D.

2.1 Reinforcement Learning

Consider a Markov Decision Process (MDP) defined as M = ⟨S, A, P, R, γ⟩, with the state space S, the action space A, the transition probability P : S × A × S → [0, 1], the reward function R : S × A → R and the discount factor γ ∈ [0, 1). An agent interacts with the MDP M by performing its policy π : S×A → [0, 1], generating the trajectories as at ∼ π(·|st), st+1 ∼ P (· | st, at) and rt = R(st, at). The value function of policy π is defined as V π (s) = Eπ [∑∞t=0 γ t rt | s0 = s]. The objective of an RL agent is to optimize its policy to maximize the expected discounted cumulative rewards J(π) = Es0∼ρ0(·)[V π (s0)], where ρ0 is the initial state distribution.

2.2 Mutual Information Optimization

Mutual Information (MI) is a widely used metric that estimates the relationship between pairs of variables in machine learning. Given two random variables X and Y , MI is a measure of dependency among them, defined as I(X; Y ) = Ep(x,y) [ p(x,y) p(x)p(y) ], with the joint probability distribution p(x, y) and the marginal probability distributions p(x), p(y). MI estimation and optimization are fundamental techniques in representation learning, with Information Bottleneck [Alemi et al., 2017] as a representative.

In most problems, it is intractable to compute MI directly since the exact distributions of interested variables are not known. Thus, approximate estimation of MI is usually resorted to, which is non-trivial especially in high-dimensional space. For the general purpose of MI optimization, lowerbound and upper-bound approximations of MI are often used as surrogate objectives [Poole et al., 2019]. Among the options, the variational approximation is a common way to compute lower bounds of MI:

The lower bound holds since the KL divergence is nonnegative. For the upper bound of MI, a recent work called Contrastive Log-ratio Upper Bound (CLUB) [Cheng et al., 2020] provides the following form:

where q(x | y) is approximated, e.g., by maximum likelihood estimation (MLE), of actual conditional distribution p(x | y) when p is not known.

3 Methodology

In this section, we propose our framework for fast adaptation, named Policy Adaptation with Decoupled Representation (PAnDR). We follow the paradigm of [Raileanu et al., 2020] and consider the offline-training-online-adaptation setting with a set of training environments Ωtrain = {Mi}Mi=1 and a set of testing environments Ωtest = {Mj}Nj=1 . Given a set of training policies Πtrain = {πk}Ok=1 , an offline experience buffer is established as D = ⋃M,O i=1,k=1 Di,k, where Di,k contains the s, a, r-trajectories generated by πk in Mi. M,N,O denote the sizes of corresponding set. In this paper, we call a set of (s, a, r) generated in an environment as the context of it, and a set of (s, a) collected by a policy as its behavior, denoted by c and b respectively.

With offline experience buffer D, PAnDR learns the representations of environment and policy through context contrast and policy recovery, respectively (Sec. 3.1). To eliminate the redundant information entangled in the two representations while retaining necessary information, we further introduce a novel approach by mutual information minimization and maximization to refine the representations (Sec. 3.2). Thereafter, we train a nonlinear neural network to approximate PDVF conditioning on both environment and policy representations and introduce the online adaptation with gradientbased policy optimization (Sec. 3.3).

3.1 Self-supervised Representation Learning

The first part of PAnDR is to learn effective representations for environment and policy from offline experiences, replying on which the environment generalization and policy adaptation are carried out. We learn these representations in a selfsupervised manner via context contrast and policy recovery, to capture the pivotal information of environment variation and policy behavior. We introduce the details below.

Learning Environment Representation via Context Contrast. By following the general idea of self-supervisedContrastive Learning [He et al., 2020], we make use of the contrast of contexts collected in different training environments to extract context embedding as environment representation. For a training environment Mi, we use the contexts sampled from Di,· = ⋃Ok=1 Di,k as anchor sample ci and positive sample c+i ; in turn, the context sampled from the offline experiences generated in other environments, i.e., Di′ ,· for any i′ ̸= i, is taken as negative sample c-i .

With the contrastive context samples described above, we train a context encoder ϕe that maps context to embedding ze by ensuring the representation of similar sample pairs (anchors and their positives) stay close to each other while dissimilar ones (anchors and their negatives) are far apart. Formally, the context encoder is optimized with InfoNCE loss similar to [He et al., 2020; Laskin et al., 2020]:

where W denotes the parameter matrix updated along with ϕe. Note that ze is differentiable with respect to ϕe; while z+ e , ze are the embeddings for positive and negative context samples calculated by a momentum averaged version of ϕe.

Compared to learning environment representation by transition-level dynamics prediction as adopted in [Raileanu et al., 2020], PAnDR captures environment variations in an integral manner and emphasizes inter-environment differences by context-based contrastive learning. Thus, the environment representation can be more effective and robust, as similarly demonstrated in [Fu et al., 2021].

Learning PolicyRepresentation via PolicyRecovery. For a training policy πk, we learn the embedding of the behavior data sampled from D·,k = ⋃Mi=1 Di,k as policy representation. As shown in the left of Fig. 2, a behavior b is fed in the behavior encoder ϕπ, producing the embedding zπ. With the purpose of making zπ contain the essential policy information, we utilize a zπ-conditioned policy decoder ϕdec π which takes as input a state s and the embedding zπ, then deterministically outputs a prediction of the actual action made by π. We call this policy recovery (or Behavioral Cloning [Grover et al., 2018]). With the offline experiences, ϕπ, ϕdec π are updated by minimizing the l2 prediction loss:

3.2 MI-based Representation Refinement

A consequent concern on the representations obtained above is the redundant and entangled information in them. This is because the representations for both environment and policy are learned from the offline experiences. The entangled information exists in both representations will negatively affect the generalization among environments and policies, thus hinders value approximation and policy adaptation in turn.

To address this problem, we propose a novel informationtheoretic approach to refine the learned representations. The key idea is 1) to eliminate the redundant information contained in environment and policy representations for the decoupleness, while 2) to ensure the essential information in their joint representation are completely retained in the two representations. This is achieved by two ways of mutual information (MI) optimization, as illustrated in the middle part of Fig. 2.

MI Minimization for Decoupleness. Firstly, for the decoupleness, we are interested in reducing the entangled information in the representations of environment and policy. To be specific, we minimize the MI between the embeddings ze and zπ that are obtained from context and behavior data sampled from the offline experiences, i.e., min I(ze; zπ). We then convert the MI minimization into minimizing the CLUB upper bound of MI, as introduced in Sec. 2.2. With the approximation conditional distribution qψ1 parameterized by ψ1, our loss function for representation decoupleness (RD) is formu-

Concretely, as to any Di,k, we use the embedding pairs of ci, bk ∼ Di,k for the samples from joint distribution p(ze, zπ); and use those of ci ∼ Di,·, bk ∼ D·,k for the marginal distributions p(ze), p(zπ). Note that ψ1 is not updated by LRD, while the approximate conditional distribution qψ1 is trained through MLE by minimizing the following loss:

where the joint distribution is sampled in the same way as in Eq. 3. The updates defined in Eq. 3 and Eq. 4 are performed iteratively: qψ1 approximates the conditional distribution with respect to the embeddings given by ϕe, ϕπ; and in turn qψ1 is used in the estimation of CLUB. We refer the readers to the original paper of CLUB [Cheng et al., 2020] for better understanding.

MIMaximization for Completeness. Secondly, we expect the information contained in the joint representation of environment and policy can be completely retained in both environment and policy representations. This is also necessary especially when the decoupleness described above is required, in case that the vital information of environment variation or policy behavior is overly reduced during MI minimization. This is achieved by maximizing the MI between the concatenation of the embeddings ze, zπ and a target environmentpolicy joint representation zb which contains complete information, i.e., maxI(ze, zπ; zb). We resort to extracting the target representation zb also from the offline experiences. Similar but distinct to the environment representation, zb is learned by context-behavior joint contrast, i.e., only the context-behavior data sampled from the same Di,k are regarded as mutually positive samples. We omit the similar formulations here and an illustration is shown in the right part of Fig. 2. Thereafter, we maximize the variational lower bound of I(ze, zπ; zb), i.e., minimize the loss function for representation completeness (RC) below:

where qψ2 is the variational distribution parameterized by ψ2, and the joint distribution are sampled by ci, bk ∼ Di,k. Note that the negative entropy term -H(zb) = Ep(zb)[log p(zb)] irrelevant to the optimization of ze, zπ is neglected in Eq. 5.

Therefore, combining the self-supervised learning losses and the MI-based refinement losses, the total loss function for the representations of environment and policy is:

where α, β are the balancing weights. The pseudocode of overall representation training is in Alg. 1 in Appendix.

3.3 Value Approximation and Policy Adaptation

So far, we have introduced the training process of environment and policy representations on the offline experiences. With the trained representations, we are ready to introduce the approximation of PDVF [Raileanu et al., 2020] and the policy optimization during online adaptation.

Value FunctionNetworkApproximation. As described in Section 1, to evaluate the values of different environmentpolicy configurations, PDVF additionally conditions on the representations of environment and policy. Distinct from the quadratic approximation adopted in [Raileanu et al., 2020], PAnDR leverages a typical multi-layer nonlinear neural network, denoted by Vθ(s, ze, zπ) with parameter θ. This setting frees the function expressivity constrained by quadratic form. We suggest that this is important because the value approximation for multiple environments and policies is much more difficult than the conventional situation (e.g., V π ). Achieving better value approximation is crucial to the following policy adaptation.

For any offline experience Di,k, we can obtain the corresponding representations ze, zπ, conditioning on which Vθ is trained by Monte Carlo method [Sutton and Barto, 1998]:

where G(s) = ∑Tt=0 γ t rt letting s0 = s is the Monte Carlo return. The pseudocode is provided in Alg. 2 in Appendix.

Policy Adaptation via Representation Gradient Ascent. After training Vθ network well, in principle we can evaluate the values for any environment-policy configuration to a certain degree. For online adaptation, we first sample a training policy from Πtrain and use it to interact with the testing environment for a few steps. Then the embedding of testing environment zon e can be inferred with the collected context. Next, starting from an arbitrary policy representation z0π (randomly initialized by default), by gradient ascent (GA), we can optimize the policy representation along the direction of increasing values of Vθ with learning rate η:

After a certain number of GA, the obtained policy representation can be used to interact with the testing environment with the help of the zπ-conditioned policy decoder ϕdec π (described in Sec. 3.1), which is trained offline to predict policy's action. We adopt the best-performing policy representation as z∗π encountered during the GA process for our adapted policy in our experiments.

It is worth noting that the GA optimization itself is zeroshot, i.e., without the overhead of interacting with the testing environment. The pseudocode is in Alg. 3 in Appendix.

4 Experiment

In our experiments, we mainly aim at investigating the following questions: (Q1) Whether PAnDR outperforms existing algorithms for fast adaptation? (Q2) How do contrastive environment representation and MI-based refinement contribute to the performance of PAnDR? (Q3) Whether nonlinear neural approximation of PDVF and gradient-based policy adaptation improve the quadratic counterpart?

4.1 Setup

We conduct experiments in four continuous control domains. All the experimental details can be found in Appendix B.

Domains and Variations. Table 1 shows the fast adaptation domains used in our experiments, where the first 3 domains are from [Raileanu et al., 2020] and the last one is from [Lee et al., 2020]. In each domain, we have a variety of environments controlled by corresponding variations. These variations determine the dynamics and reward functions of environment. Therefore, fast adaptation is expected to be achieved against these changes. The number of training and testing environments are provided in Table 1.

Offline Experiences and Adaptation Budget. We follow the same data setting as [Raileanu et al., 2020]. For offline experiences, a policy is optimized by PPO [Schulman et al., 2017] in each training environments. The obtained policies are training policies that serve as collectors. In general, such collector policies can be obtained from other ways. For each domain, 50-episode interaction trajectories are collected by each combination of training environment and policy, e.g., 15 ∗ 15 ∗ 50 episodes of experience for Ant-Wind. For online adaptation, 50 steps of interaction are collected in the testing environment by one of the training policies. For PAnDR, 100 steps of GA is used by default.

Evaluation Criteria. The main evaluation criterion is adaptation performance, calculated as the undiscounted cumulative rewards of the policy obtained after online adaptation. Moreover, for PAnDR, we analyze the evolvement of adaptation performance against the number of GA adaptation steps (in Sec. 4.3). The results reported in our experiments are means and error bars of a standard deviation calculated with 7 independent trials.

4.2 Performance Comparison

In this section, we compare the adaptation performance between PAnDR and representative existing algorithms to answer Q1 within the domains introduced above.

Baselines. We consider the following baselines. 1) PPO: We use the PPO policy which is trained online with 3e6 steps in testing environments as a powerful baseline. 2) MAML [Finn et al., 2017]: We use MAML as a representative of gradient-based Meta RL. Note MAML (as well as most Meta RL algorithms) performs online meta training. This means in principle it is not applicable to the offline training experiences. 3) PDVF [Raileanu et al., 2020]: Since we propose PAnDR based on the same paradigm of PDVF, we use PDVF as a closely related baseline. Moreover, to our knowledge, PDVF is the state-of-the-art fast adaptation algorithm in offline-training-online-adaptation setting. Note that in this section, we use PDVF to denote the method name of [Raileanu et al., 2020] rather than the extended value function. For all above baselines, we use the official source codes provided by PDVF paper [Raileanu et al., 2020]. For PAnDR, we also implement upon these source codes with almost the same configuration of hyperparameters and structures for a valid comparison.

(A1) Results. As shown in Fig. 3, for the simple Spaceship domain, the overall performance of PAnDR is significantly better than PDVF and slightly outperfroms MAML. Further in the three more difficult MuJoCo domains, the leading edge of PAnDR is much more evident and consistent compared with the two baseline methods, while MAML loses its competitiveness. We consider that this is because the preferable representations obtained via self-supervised learning and the MI-based refinement are more effective and robust than those of PDVF. Moreover, with less-constrained approximation, the more accurate value estimation achieved improves the optimality of policy obtained during adaptation.

It is worth noting that, perhaps somewhat surprisingly, PAnDR shows competitive and even better performance when comparing with the powerful PPO baseline trained in the testing environments. Instead of requiring millions of interactions to optimize a policy in the testing environments, PAnDR only needs limited interacting experiences for fast adaptation. This further reveals the superiority and potential of PAnDR.

4.3 Ablation Study and Analysis

In this section, we give the answers to Q2 and Q3. We provide main results below, and the complete results along with detailed analysis are given in Appendix A.

(A2-I) MI-based Representation Refinement. To verify the contribution of the MI minimizing and maximizing in representation refinement of PAnDR, we compare PAnDR with its several variants including: 1) PAnDR without MAX MI: removes MI maximizing refinement; 2) PAnDR without MIN MI: removes MI minimizing refinement; 3) PAnDR without MI: removes both MI maximizing and minimizing refinements. As shown in Fig.4, both MI minimizing and maximizing positively contribute to the improvement of PAnDR's overall performance. It demonstrates that both MI-based refinement components are indispensable in PAnDR.

(A2-II) Context Contrast for Environment Representation. Fig.5, presents the efficacy of context contrast in comparison with dynamics prediction upon both types of Vθ in PDVF and PAnDR. The associated variants are shown in Table 2. As in Fig.5, PDVF with CER has superior performance than original PDVF while PAnDR without MI outperforms PAnDR with NA-GA. Thus, context contrast is more favorable in both cases of Vθ compared with dynamics prediction.

(A3-I) Vθ Approximation & GA Adaptation. With the variants in Table 2, we compares: PAnDR v.s. PAnDR with QA, and PDVF v.s. PDVF with NA-GA. The results in Fig.5 verify that there is a clear advantage of using NN nonlinear approximation of Vθ along with GA optimization over the quadratic counterpart.

(A3-II) Adaptation Performance against GA Steps. Since there is no longer a closed-form solution to optimization available for Vθ with NN nonlinear approximation, we investigate how consecutive steps of GA affect the adpatation performance of PAnDR. The results for various choices of GA step number are provided in Table 3-8 in Appendix A.3. Moreover, we provide detailed analysis for useful insights.

5 Conclusion and Discussion

This paper focus on fast adaptation problem with only offline experiences on training environments available. Following the paradigm of [Raileanu et al., 2020], we propose a new algorithm, PAnDR, that consists of effective environment and policy representations and better policy optimization during online adaptation. In essence, policy adaptation in PAnDR (and PDVF) is a process of policy representation optimization and decoding. Both the two parts are not well studied at present. Optimization in policy representation space is nontrivial as there is a lack of knowledge on the landscape of the space; while the decoding of such latent representation can be invalid without carefully handling the uncertainty [Notin et al., 2021]. We consider that they are significant and there is much room for improvement. In addition, advances in environment and policy representation learning are expected to further improve the performance of PAnDR.

Acknowledgments

This work is supported by the National Science Fund for Distinguished Young Scholarship of China (No. 62025602), the National Key R&D Program of China (Nos. 2018AAA0100900), the National Natural Science Foundation of China (Nos. U1803263, 62106172 and 11931015), the Key Technology Research and Development Program of Science and Technology-Scientific and Technological Innovation Team of Shaanxi Province (Grant No. 2020TD-013), the XPLORER PRIZE, the New Generation of Artificial Intelligence Science and Technology Major Project of Tianjin (Grant No.: 19ZXZNGX00010), and the Science and

• Technology on Information Systems Engineering Laboratory (Grant No. WDZC20205250407).