Title: State Regularized Policy Optimization on Data with Dynamics Shift
Venue: NeurIPS 2023
Accepted: yes
Keywords: Reinforcement Learning

Dynamics Shift

Stationary State Distribution

Offline RL

Off-Policy RL
Abstract: In many real-world scenarios, Reinforcement Learning (RL) algorithms are trained on data with dynamics shift, i.e., with different underlying environment dynamics. A majority of current methods address such issue by training context encoders to identify environment parameters. Data with dynamics shift are separated according to their environment parameters to train the corresponding policy.
However, these methods can be sample inefficient as data are used \textit{ad hoc}, and policies trained for one dynamics cannot benefit from data collected in all other environments with different dynamics. In this paper, we find that in many environments with similar structures and different dynamics, optimal policies have similar stationary state distributions. We exploit such property and learn the stationary state distribution from data with dynamics shift for efficient data reuse. Such distribution is used to regularize the policy trained in a new environment, leading to the SRPO (\textbf{S}tate \textbf{R}egularized \textbf{P}olicy \textbf{O}ptimization) algorithm. 
To conduct theoretical analyses, the intuition of similar environment structures is characterized by the notion of homomorphous MDPs. We then demonstrate a lower-bound performance guarantee on policies regularized by the stationary state distribution. In practice, SRPO can be an add-on module to context-based algorithms in both online and offline RL settings.
Experimental results show that SRPO can make several context-based algorithms far more data efficient and significantly improve their overall performance.

State Regularized Policy Optimization on Data with Dynamics Shift

Abstract

In many real-world scenarios, Reinforcement Learning (RL) algorithms are trained on data with dynamics shift, i.e., with different underlying environment dynamics. A majority of current methods address such issue by training context encoders to identify environment parameters. Data with dynamics shift are separated according to their environment parameters to train the corresponding policy. However, these methods can be sample inefficient as data are used ad hoc , and policies trained for one dynamics cannot benefit from data collected in all other environments with different dynamics. In this paper, we find that in many environments with similar structures and different dynamics, optimal policies have similar stationary state distributions. We exploit such property and learn the stationary state distribution from data with dynamics shift for efficient data reuse. Such distribution is used to regularize the policy trained in a new environment, leading to the SRPO ( S tate R egularized P olicy O ptimization) algorithm. To conduct theoretical analyses, the intuition of similar environment structures is characterized by the notion of homomorphous MDPs. We then demonstrate a lower-bound performance guarantee on policies regularized by the stationary state distribution. In practice, SRPO can be an add-on module to context-based algorithms in both online and offline RL settings. Experimental results show that SRPO can make several context-based algorithms far more data efficient and significantly improve their overall performance.

1 Introduction

Reinforcement Learning (RL) has achieved great success in solving challenging sequential decisionmaking problems [1, 2]. Unfortunately, existing RL methods usually assume that agents are trained and evaluated in exactly the same environment, which is often not the case in real-world applications where environment dynamics can vary a lot. For example, the recommendation engine of social apps may need to deal with time-varying and heterogeneous user preferences [3, 4]. A robot arm may operate in different scenarios with different joint frictions and medium densities [5]. In these cases, the agent has to work with the trajectory data from different environment dynamics, i.e., data with dynamics shift , which will bias the learning process and lead to poor performance. In fact, some empirical studies [6, 5] demonstrate that general RL algorithms [7, 8] can easily be misled by different environment dynamics and fail to train a good policy.

In recent years, considerable research efforts have been devoted to addressing the dynamics shift and learning generalizable policies for environments with changing dynamics. One common practice is to train a context encoder [9, 10, 11] to associate the environment dynamics with a latent variable. The policy is then trained with the latent variable as an additional input [12]. One issue with this

practice is that policies conditioned on a specific latent variable can only learn from data collected in the environment corresponding to that latent variable. In other words, data with different dynamics are used in an ad hoc manner. The generalizability of context encoders relies on the expressive power of neural networks. However, neural networks are prone to overfit and behave poorly when extrapolating. As an example, we benchmarked CaDM [9], which is one of the context-based algorithms, under Ant environments with different gravities and display the results in Fig. 1. Although it can outperform PPO [7] due to its adaptability from context encoders, CaDM fails to constantly improve its performance with more data from different environment dynamics. To mitigate the problem of inefficient data use, there are some attempts that leverage Importance Sampling (IS) [13, 5, 14]. Given the dynamics of the target environment, samples from the source environments are assigned with larger importance weights if they are more likely to happen in the target environment and vice versa. Compared to training context

1

3

5

Number of Different Dynamics

0

1

2

3

4

Return

×10 2

Performance in Ant Environment

Algorithm

PPO

CaDM+SRPO

encoders, IS-based methods manage to proactively exploit the data from other dynamics. However, such methods require prior knowledge about the dynamics of the target environment. Also, it is notoriously hard to balance the bias and variance when calculating the IS weights.

This paper proposes a new RL paradigm that can explicitly leverage data with dynamics shift. It is also free of the aforementioned drawbacks of IS-based methods. We find that the stationary state distribution induced by optimal policies (later termed optimal state distribution) is similar across a set of environments with similar structures and different environment dynamics. For example, given heterogeneous preferences of users, a video recommendation system may choose different videos to recommend, but the optimal states are the same: users keep pressing the 'like' or 'save' button and continue watching for a long time. More concretely, the optimal state distribution in one environment dynamics can be informative for training policies in all other different dynamics. We therefore propose a constrained policy optimization (CPO) [15] formulation that requires the policy not only to optimize the cumulative return, but also to generate a stationary state distribution close to the optimal state distribution. By relating optimality to high-reward states [16], we are able to approximate the optimal state distribution from trajectory data regardless of the underlying dynamics, providing a unified and efficient approach to exploiting these data.

Summarizing these ideas, we propose the SRPO ( S tate R egularized P olicy O ptimization) algorithm. SRPO works as an add-on module in both online and offline context-based RL algorithms such as CaDM [9] and MAPLE [12] to increase their sample efficiency, leading to the CaDM+SRPO and MAPLE+SRPO algorithms. We provide a lower-bound performance guarantee on policies in one dynamics regularized by the optimal state distribution in other dynamics. This theoretically demonstrates the effectiveness of the SRPO algorithm in using data with dynamics shift. Empirical results in both online and offline settings show that SRPO can significantly improve both the data efficiency and the overall performance of several state-of-the-art context-based RL algorithms. We also perform ablation studies to demonstrate the effectiveness of each component in the SRPO algorithm.

2 Backgroud

2.1 Preliminaries

A Markov Decision Process (MDP) can be defined by a tuple ( S , A , T, r, γ, ρ 0 ) , where S is the state space, A is the bounded action space with actions a ∈ ( -1 , 1) , T ( s ′ | s, a ) ∈ [0 , 1] and r ( s, a, s ′ ) ∈ [ -R max , R max ] are the transition and reward functions. γ ∈ (0 , 1) is the discount factor and ρ 0 ( s ) is the initial state distribution. In MDPs with deterministic transitions, we denote T ( s, a ) as the transition function with a slight abuse of notation, and ( T, ε ) as { T ′ | | T ( s, a ) -T ′ ( s, a ) | < ε, ∀ s ∈ S , a ∈ A} which is the ε -neighbourhood of T . RL aims at maximizing the accumulated return of policy π :

η T ( π )= E π,T [ ∞ ∑ t =0 γ t r ( s t , a t ) ] , where the expectation is computed with s 0 ∼ ρ 0 , a t ∼ π ( ·| s t ) , and s t +1 ∼ T ( ·| s t , a t ) . The optimal policy π ∗ is defined as π ∗ T = arg max π η T ( π ) . In an MDP with a policy π , the Q-value Q π T ( s, a ) denotes the expected return after taking action a at state s : Q π T ( s, a )= E π,T [ ∞ ∑ t =0 γ t r ( s t , a t ) | s 0 = s, a 0 = a ] . The value function is defined as V π T ( s ) = E a ∼ π ( ·| s ) Q π T ( s, a ) with V ∗ T ( s ) being the shorthand for V π ∗ T T ( s ) . It satisfies the optimal Bellman Equation V ∗ T ( s ) = max a r ( s, a ) + γ E s ′ ∼ T ( ·| s,a ) V ∗ T ( s ′ ) . We can also define the stationary state distribution (also known as state occupation function) as d π T ( s ) := (1 -γ ) ∑ ∞ t =0 γ t P T ( s t = s | π ) with d ∗ ( s ) being the shorthand for d π ∗ T ( s ) .

T T

̸

The Hidden Parameter Markov Decision Process (HiP-MDP) captures a class of MDPs with different transition functions and the same reward function by introducing a set of hidden parameters. Specifically, an HiP-MDP is defined by a tuple ( S , A , Θ , T, r, γ, ρ 0 ) , where Θ is the space of hidden parameters. The transition function T θ ( s ′ | s, a, θ ) is parameterized not only by states and actions, but also by a hidden parameter θ sampled from Θ . The action gap of an HiP-MDP is defined as ∆ = min θ ∈ Θ min s ∈S min a = π ∗ ( s ) V ∗ T θ ( s ) -Q ∗ T θ ( s, a ) , which reflects the minimum gap between an optimal action and all other sub-optimal actions.

2.2 Related Work

MDPs with Different Dynamics The setting of HiP-MDP [17] was proposed to model a set of variations in the environment dynamics. The problem is intensively investigated in recent years [18] and these researches fall into three categories, i.e., encoder-based, Important Sampling (IS)-based and meta-RL based algorithms. Encoder-based methods extract the hidden parameters from trajectories with variational inference [19] or auxiliary loss [9]. These hidden parameters are used as inputs to the transition function [9] or policy network [20, 12]. Unfortunately, these methods train dynamicsspecific policies from the trajectory data of each hidden parameter independently, which leads to poor sample efficiency. Instead, our method uses the data from all dynamics to learn an optimal state distribution that facilitates the policy learning. IS-based methods compute the importance ratio between transition probabilities under different dynamics and modify the replay buffer [13, 14, 5] according to the transition probabilities in the test environments, which is often not available in real-world scenarios. Finally, meta-RL algorithms [21, 22] can adapt to environments with new dynamics through fine-tuning on a small amount of data from the test environment. In contrast, our method can be directly applied to new environment dynamics by making a zero-shot transfer.

Behavior Regularized Methods The idea of constrained policy optimization (CPO) [15] is widely used in RL. Most researches focus on behavior regularized methods, i.e., adding policy constraints based on another policy distribution, as shown in the following optimization problem:

where D is the replay buffer, ˆ π is the regularizing policy, and ˆ D is a certain distance measure. Maximum-Entropy RL [23, 8] can be considered as CPO with a uniform policy distribution. The sparse action tasks [24] can be solved by CPO with a sparse policy distribution. Besides, many offline RL algorithms [25, 26, 27] are based on the idea of constraining the current policy distribution to be close to the dataset's policy distribution. However, data sampled from environments with different dynamics can have distinct optimal policies, as illustrated in Sec. 3.1. In such cases, ˆ π in Eq. (1) may include policies that do not match the current environment, and can therefore be misleading. So behavior regularization in Eq. (1) would fail on data with dynamics shift. Differently, our proposed method is based on state regularization , which is more suitable when learning from data with dynamics shift.

Leveraging stationary state distributions The stationary state distribution d π T ( s ) of policy π and dynamics T is an important feature that can measure the differences in policies and transition functions. It has already been exploited in many researches. In Off-Policy RL, Islam et. al [28] estimates the stationary state distributions of both the current policy and the mixed buffer policy. It

then computes the off-policy policy gradient with the constraint that the two distributions should be close. Some Off-Policy Policy Evaluation (OPE) algorithms [29, 30] use the steady-state property of stationary Markov processes to estimate the stationary state distributions. In Imitation Learning (IL), state-only IL algorithms [31, 32] requires the stationary state distribution of the current policy to be close to that of the expert policy. In Inverse RL (IRL), [33] learns a stationary reward function by computing the gradient of the distance between agent and expert state distribution w.r.t. reward parameters. In Offline RL, [34] requires the stationary state distribution of the learning policy and the behavior policy to be close and perform conservative updates. The use of such distributions in our paper is similar to some researches on sim-to-real [35, 36]. They propose to match the next state distribution in the imperfect simulator and the real environment with inverse dynamics model. They implicitly relies on the idea that the same state distribution should generate similar returns in environments with different dynamics. We formulate the idea in this paper with theorems and quantitatively analyse such similarity in various conditions.

3 State Regularized Policy Optimization

In this section, we first give motivating examples on why the optimal state distribution in one environment dynamics can be informative in all other different dynamics. A constrained policy optimization formulation is then proposed in Sec. 3.2 based on the optimal state distribution. Solving this optimization problem gives rise to our State Regularized Policy Optimization (SRPO) algorithm that can leverage data with dynamics shift to improve the policy performance.

3.1 Motivating Example

The key intuition behind SRPO is that the optimal state distribution is similar across environments. Consider an example of the Inverted Pendulum environment in Fig. 2. We train two policies in the environments with gravities of 5 and 10 until convergence. Then the kernel density estimation [37] technique is employed to estimate the state and action density of the data collected by the two policies in different areas. It can be observed from the figure that collected data have the same high state density region with low pendulum speed and small pendulum angle, while the action distribution has different density peaks. It demonstrates that the state distribution of data generated by the optimal policy can be similar regardless of the environment dynamics, and therefore can serve as a reference distribution to regularize the training policy in environments with new dynamics. More demonstrating examples can be found at Appendix B.3.

3.2 State Regularized Policy Optimization

Based on the intuition of informative optimal state distribution, we develop a novel technique that regulates RL algorithms to generate a stationary state distribution that is close to the optimal one. Specifically, we propose the following constrained policy optimization formulation:

0.0

0.8

1.6

2.4

3.1

Angle

-4.0

-2.0

0.0

2.0

4.0

Speed

State Density with Gravity 5

0.0

0.8

1.6

2.4

3.1

Angle

-4.0

-2.0

0.0

2.0

4.0

State Density with Gravity 10

25

20

15

10

5

0

-2.0

-1.0

0.0

1.0

2.0

Torque

Action Density with Gravity 5

-2.0

-1.0

0.0

1.0

2.0

Torque

Action Density with Gravity 10

5

0

where ζ ( s ) is the optimal state distribution in other environment dynamics. By introducing the stationary state distribution, the optimization problem defined in Eq. (2) extends the regularization of in-distribution data to data with distribution shift 1 .

We solve Eq. (2) by casting it to the following unconstrained optimization problem via Lagrange multipliers 2 :

where λ > 0 is the Lagrangian Multiplier. It is noteworthy that in addition to the multiplier term, the only difference of Eq. (3) and the reward-maximization objective of RL is that the logarithm of probability density ratio λ log ζ ( s t ) d π ( s t ) is added to the reward term r ( s t , a t ) . Therefore, one can easily apply our scheme to a wide range of RL algorithms by augmenting the reward function with the density ratio.

3.3 Data-based Surrogate of the Density Ratio

The main challenge in solving Eq. (3) is to compute the density ratio ζ ( s ) d π ( s ) because obtaining the optimal state distribution during online training or given suboptimal offline dataset is infeasible. Also, d π ( s ) is intractable if the state space is continuous. Motivated by recent advances in adversarial training [38, 39], we propose a sample-based surrogate for the density ratio ζ ( s ) d π ( s ) .

Proposition 3.1. In a GAN, when the real data distribution is ζ ( s ) and the generated data distribution is d π ( s ) , the output of the discriminator D ( s ) follows

We discuss the relation of this sample-based surrogate with f-divergences and Off-Policy RL in Appendix A.3. To train the discriminator D ( s ) , we need to generate samples that is close to the optimal state distribution ζ ( s ) and away from d π ( s ) , which is sort of the average state distribution. Motivated by [16], we model state optimality by a variable O t . As shown in Fig. 3, we regard the state s t in MDP as a hidden state in a Hidden Markov Model (HMM), and introduce the binary observation state O t . O t = 1 denotes that s t is the optimal state at timestep t . The observation model is given by

s

1

a

1

O

1

s

2

a

2

O

2

s

3

a

3

O

3

s

4

O

4

We can therefore compute the state density ratio ζ ( s ) d π ( s ) as

where the second equation follows the Bayes' law. The last term is related to the forward probability α t ( s t ) = p ( O 0: t -1 | s t , π ) and backward probability β t ( s t ) = p ( O t : ∞ | s t , π ) in the HMM. We discuss in Appendix A.2 that β t ( s t ) is positively related to a soft version of MDP's state value V π ( s t ) . Also, α t ( s t ) makes a little influence on the overall density ratio. Therefore, the input s will be more likely to be sampled from distribution ζ ( s ) if it has a higher state value V ( s ) than average. With this idea, we are able to build training samples for the discriminator D ( s ) .

Algorithm 1 The workflow of SRPO on top of MAPLE [12].

1: Input: ϕ φ as a context encoder parameterized by φ ; Adaptable policy network π θ parameterized by θ ; Adaptable value network V ψ parameterized by ψ ; Offline dataset D off ; Rollout horizon H ; State partition ratio ρ ; State discriminator D δ parameterized by δ ; Regularization coefficient λ .

2: for 1, 2, 3, . . . do

3: for t = 1 , 2 , . . . , H do

4: Sample z t from ϕ φ ( z | s t , a t -1 , z t -1 ) and then sample a t from π θ ( a | s t , z t ) .

5: Rollout and get transition data ( s t +1 , r t , d t +1 , s t , a t , z t ) . Then add it to D rollout .

6: end for

7: Update the context encoder ϕ φ according to MAPLE.

8: Sample a batch D batch from D off and D rollout and rank them by their state-values estimated by V ψ ; Add ρ |D batch | states with higher state-values to D real and the others to D fake .

9: Train the discriminator D δ with nll loss.

10: For one-step transition ( s t +1 , r t , d t +1 , s t , a t , z t ) in D batch, update r t with r t + λ D δ ( s t ) 1 -D δ ( s t ) .

11: Use the updated D batch and SAC to update the policy and value network parameters θ and ψ . 12: end for

3.4 Practical Algorithm

Summarizing the previous derivations, we obtain a practical reward regularization algorithm, termed as SRPO ( S tate R egularized P olicy O ptimization) to leverage data with dynamics shift. We select the MAPLE [12] algorithm, which is one of the SOTA algorithms in context-based Offline RL, as the base algorithm. The detailed procedure of MAPLE+SRPO is shown in Alg. 1. After preparing the dataset in a model-based Offline RL style [40, 12], we sample a batch of data from the dataset, obtain a portion of ρ states with higher rewards and add them to the dataset D real . D fake is similarly generated by states with lower rewards (line 10). We set ρ = 0 . 5 in offline experiments with medium-expert level of data. ρ = 0 . 2 is set in all other experiments. Then a classifier discriminating data from the two datasets is trained (line 11). It estimates the logarithm of the state density ratio λ log ζ ( s ) d π ( s ) , which is added to the reward r t (line 12). λ is regarded as a hyperparameter with values 0 . 1 or 0 . 3 . The effect of λ is investigated in Sec. 5.3. The procedure of the online algorithm CaDM [9]+SRPO is similar to MAPLE+SRPO, where the datasets D real and D fake are built with data from the replay buffer, rather than the offline dataset.

4 Theoretical Analysis

In this section, we analyze some properties of MDPs with different dynamics and provide theoretical justifications for the SRPO algorithm in Sec. 3. The notations are introduced in Sec. 2.1 and proofs can be found in Appendix A.4. We first show in Thm. 4.2 that the performance of a policy can be lower-bounded when its stationary state distribution is close to a certain optimal state distribution. In accordance with the intuition in Sec. 3.1, it is also demonstrated in Thm. 4.3 that optimal policies can induce the same stationary state distribution in different dynamics under mild assumptions. We start the analysis with the definition of homomorphous MDPs.

Definition 4.1 (homomorphous MDPs) . In an HiP-MDP ( S , A , Θ , T, r, γ, ρ 0 ) , consider hidden parameters θ 1 , θ 2 ∈ Θ . Let T i ( s ′ | s, a ) = T ( s ′ | s, a, θ i ) , ∀ ( s, a, s ′ ) ∈ S × A × S , i = 1 , 2 . If ∑ a ∈A T 1 ( s ′ | s, a ) > 0 ⇔ ∑ a ∈A T 2 ( s ′ | s, a ) > 0 for all s, s ′ ∈ S , MDPs ( S , A , T 1 , r, γ, ρ 0 ) and ( S , A , T 2 , r, γ, ρ 0 ) are referred to as homomorphous MDPs.

In this definition, ∑ a ∈A T ( s ′ | s, a ) > 0 means state s ′ can be reached from s , so the equivalence of non-zero transition probabilities refers to the same reachability from s to s ′ . Such condition holds in a wide range of MDPs differing only in environment parameters. For example, pendulums with different lengths can all reach the upright state from an off-center state, with longer pendulums exerting a larger force. Apart from the homomorphous property, we also require the reward and dynamics functions of MDPs to have Lipschitz properties. We assume reward function r ( s, a, s ′ ) w.r.t. the action a is λ 1 -Lipschitz and the dynamics function T ( s, a ) w.r.t. the action a is λ 2 -inverse Lipschitz. Discussions on these Lipschitz properties can be found in Appendix A.5.

With these preliminaries, we first analyze the discrepancy of accumulated returns of two policies with similar stationary state distributions. The analysis is related to our SRPO algorithm in that the state regularized policy optimization formulation in Eq. (2) also constrains the learning policy to have a similar stationary state distribution with the optimal policy. Specifically, we derive a theorem as follows.

Theorem 4.2. Consider two homomorphous MDPs with dynamics T and T ′ . If T ′ ∈ ( T, ε m ) , for all learning policy ˆ π such that D KL ( d ˆ π T ( · ) ∥ d ∗ T ′ ( · )) ⩽ ε s , we have

The theorem implies that if a policy ˆ π has a similar stationary state distribution with the optimal policy in one MDP M , ˆ π will have a lower-bound performance guarantee in all MDPs that are homomorphous with the MDP M . Therefore, the learning policy can benefit from the state regularized policy optimization in Sec. 3.2.

More specifically, Eq. (7) shows that the gap in accumulated return of ˆ π and π ∗ T is related to the dynamics shift ε m , the KL-Divergence of two stationary state distributions ε s , and the effective planning horizon 1 1 -γ . With respect to the dynamics shift ε m , it is related to a 'uniform' constraint on the dynamics T ′ . We further show in Appendix A.4 that constraining the dynamics shift on a certain state-action pair is enough to derive Eq. (7). Unlike the dynamics shift ε m that is determined by a pre-defined RL task, the discrepancy between stationary state distributions ε s is determined by the learning policy ˆ π and can be optimized during training to obtain a better performance lower-bound. We also discuss in Appendix A.5 how tight Eq. (7) is in terms of the effective planning horizon 1 1 -γ , compared with some similar performance bounds.

With an additional assumption on the action gap ∆ (defined in Sec. 2.1), we further demonstrate that the optimal policy of two homomorphous MDPs can have the same stationary state distribution, which verifies the intuition in Sec. 3.1.

Theorem 4.3. Consider two homomorphous MDPs with dynamics T and T ′ . If T ′ ∈ ( T, ε m ) and the action gap ∆ follows ∆ > (2 -γ ) λ 1 λ 2 ε m 1 -γ , for all s ∈ S we have d ∗ T ( s ) = d ∗ T ′ ( s ) .

The assumption is mild and holds in many scenarios. For example, in autonomous driving it can be very dangerous to deviate from the optimal policy. Such suboptimal actions have low rewards, leading to a large action gap ∆ . In recommendation tasks we are hardly concerned with what items we recommend (the action), as long as the recommendation outcome (the state), i.e., the users' experiences are good enough, leading to a small λ 1 . The condition of large enough action gap holds in these situations.

5 Experiments

In this section, we conduct experiments to investigate the following questions: (1) Can SRPO leverage data with distribution shift and outperform current SOTA algorithms in the setting of HiP-MDP, in both online and offline RL? (2) How does each component of SRPO (e.g., use state regularization rather than behavior regularization) contribute to SRPO's performance? To answer question (1), we use the MuJoCo simulator [41] and generate environments with different transition functions. We train the CaDM+SRPO and the MAPLE+SRPO algorithm proposed in Sec. 3.4 and make comparative analysis with baseline algorithms. To answer question (2), we do ablation studies to examine the role of different modules in SRPO. We also examine how the discriminator D δ works in complex environments and the effect of regularizing with state distributions in different performance levels.

5.1 Experiment Setup

We alter the simulator gravity to generate different dynamics in online experiments. Possible values of gravity are {1.0}, {0.7,1.0,1.3}, and {0.4,0.7,1.0,1.3,1.6} in experiments with 1, 3, and 5 kinds of different dynamics, respectively. When the simulator resets, the gravity is uniformly sampled from the set of all possible values. The number of training steps is in proportion to the number of environment parameters. Therefore, the agent has access to the same amount of training data on a

0.0

0.3

0.6

0.9

1.2

1.5

1.4

1.2

1.0

0.8

0.6

0.4

×10 3

Pendulum (3 Different Dynamics)

0

1

2

3

4

5

6

0.0

0.5

1.0

1.5

2.0

2.5

×10 2

Ant (3 Different Dynamics)

0

1

2

3

4

5

6

0.0

0.5

1.0

1.5

2.0

2.5

3.0

×10 3

HalfCheetah (3 Different Dynamics)

0

1

2

3

4

5

6

0.0

0.2

0.4

0.6

0.8

1.0

1.2

×10 5

Humanoid (3 Different Dynamics)

0.0

0.5

1.0

1.5

2.0

2.5

Timesteps(M)

1.4

1.2

1.0

0.8

0.6

0.4

×10 3

Pendulum (5 Different Dynamics)

0

2

4

6

8

10

Timesteps(M)

0

1

2

3

4

5 ×10 2

×10 2

Ant (5 Different Dynamics)

0

2

4

6

8

10

Timesteps(M)

0

1

2

3

4 ×10 3

×10 3

HalfCheetah (5 Different Dynamics)

0

2

4

6

8

10

Timesteps(M)

0.0

0.2

0.4

0.6

0.8

1.0

1.2

×10 4

×10 5

Humanoid (5 Different Dynamics)

PPO

CaDM+SRPO

certain value of simulator gravity. We also consider the shift of medium density and body mass in offline experiments to show SRPO's robustness to different forms of dynamics shift.

To perform comparative analysis, we choose CaDM [9] and PPO [7] as baseline algorithms in online experiments. In offline experiments, DARA [5] also exploits large amount of data with dynamics shift. Its algorithm relies on Importance Sampling and will be used as a baseline method. Apart from that, we choose MOPO [40], MAPLE [12] and CQL [26] as baseline methods. More information on the setup of experiments is shown in Appendix B.1.

5.2 Results

Online Experiments The results of online experiments are shown in Fig. 4. With the context encoder and conditional policy, CaDM is able to outperform PPO in all environments. However, it fails to take advantage of the increase in the amount of data with dynamics shift. Its performance with 5 different dynamics is lower than that with 3 dynamics. In contrast, our proposed SRPO

algorithm leads to better performance on top of CaDM in accordance with more training data. It significantly outperforms the original CaDM algorithm in environments with 5 different dynamics. The performance comparison in the Pendulum environment is also in accordance with the motivating example in Sec. 3.1. More results of online experiments are shown in Appendix B.2.

Offline Experiments The results of offline experiments are shown in Tab. 1. The column of 'CQL Single' refers to the evaluation score in the CQL [26] paper, where the policy is with data from a single static environment. Without the mechanism of context-based encoders, GAIL [43], CQL and MOPO [40] cannot handle data with distribution shift and show a performance drop. MAPLE [12] and MAPLE+DARA [5] only achieve marginal performance improvement with respect to CQL single. On the other hand, MAPLE+SRPO shows significant performance improvement over CQL single, which means that SRPO can efficiently leverage the additional data with dynamics shift to facilitate policy training. The MAPLE+SRPO algorithm also has a 15% higher evaluation score than MAPLE, achieving the best performance in 8 out of 12 tasks. Apart from MAPLE, the meta-RL algorithm PEARL [44] also has an context encoder for fast adaptation. We compare PEARL with PEARL+SRPO and leave the results in Appendix. B.3.

5.3 Analysis

Ablations We conduct ablation studies in offline environments to analyze the role of each algorithm component in SRPO. The results are shown in Tab. 2. We first investigate the outcome of regularizing with state-action distribution rather than state distribution in Eq. (2). The resulting policy has a lower evaluation score on average than policies trained with the original SRPO algorithm in all environments. This is because environments with different dynamics do not have a similar optimal policy. The action distribution in the mixed dataset can be misleading when training new policies. According to Sec. 3.3, SRPO trains a classifier to discriminate states with higher values from lower values. We also train another classifier discriminating a random binary partition of states. Ablation results show a huge performance drop, which verifies the effectiveness of the classification-based surrogate mechanism. We also evaluate MAPLE+SRPO with a fixed value of the hyperparameter λ . λ = 0 . 1 is more suitable for Walker2d and HalfCheetah environments, while in the Hopper environment λ = 0 . 3 is better. This is in accordance with previous analysis that Hopper agents can benefit more from regularizing with the stationary state distribution.

Effectiveness of Discriminators We first train a discriminator D δ according to Alg. 1. Then a set of states is sampled from the D4RL [42] dataset and classified into two sets according to the output of D δ . The average values of the two sets of states are compared in Fig. 5. As shown in the figure, states with higher D δ outputs also have higher values in all three environments. It means that the trained discriminator D δ can successfully identify states with high values from those with low values. Therefore, its output can be a good surrogate for the density ratio ζ ( s ) d π ( s ) in Sec. 3.

Effectiveness of the Regularization We also study the effect of policy regularization with different performance levels of stationary state distributions. Random, medium and expert policies in the original Hopper environment are

Walker2d

Hopper

HalfCheetah

Envrironment

0.4

0.6

0.8

1.0

1.2

Normalized Values

State Values With Different D Output

High

Low

Normalized Value Gaps

1.0

0.8

0.6

0.4

0.2

Value Gaps of Different

Random

Medium

Policy Type

used to estimate the stationary state distributions, which regularize the learning policies in a new environment with different dynamics. The results are shown in Tab. 3, where the expert policy is the most effective in regularizing. This verifies the practice in Sec. 3.2 and the theoretical analysis.

6 Conclusion and Discussion

In this work, we focus on the problem of leveraging data with dynamics shift to efficiently train RL agents. Based on the intuition that optimal policies can lead to similar stationary state distributions, we give a constrained optimization formulation that regards the state distribution as a regularizer. After discussions on a sample-based surrogate, we propose the SRPO algorithm which can be an add-on module to context-based algorithms and improve their sample efficiency. The resulting CaDM+SRPO and MAPLE+SRPO algorithms show superior performance when learning on data sampled from environments with different dynamics. Theoretical analyses are also given to analyze some properties of MDPs with different dynamics. They provide justifications for the intuition of the dynamics-invariant state distribution, as well as the constrained policy optimization formulation.

Limitations and Future work The theoretical analyses of this work requires the assumption of homomorphous MDPs, i.e., the same state reachability in different MDPs. It would be interesting to discuss whether similar conclusions on the stationary state distribution can be derived without such assumption.

Acknowledgements We thank Wanqi Xue and Yanchen Deng for helpful discussions. This research is supported by the National Research Foundation, Singapore under its Industry Alignment Fund - Prepositioning (IAF-PP) Funding Initiative and Ministry of Education, Singapore, under its Academic Research Fund Tier 1 (RG13/22). Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Singapore.

References

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin A. Riedmiller. Playing atari with deep reinforcement learning. CoRR , abs/1312.5602, 2013.

David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Vedavyas Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy P. Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of go with deep neural networks and tree search. Nat. , 529(7587):484-489, 2016.

Wanqi Xue, Qingpeng Cai, Zhenghai Xue, Shuo Sun, Shuchang Liu, Dong Zheng, Peng Jiang, and Bo An. PrefRec: Preference-based recommender systems for reinforcing long-term user engagement. CoRR , abs/2212.02779, 2022.

Zhenghai Xue, Qingpeng Cai, Tianyou Zuo, Bin Yang, Lantao Hu, Peng Jiang, and Bo An. AdaRec: Adaptive sequential recommendation for reinforcing long-term user engagement. CoRR , abs/2310.03984, 2023.

Jinxin Liu, Hongyin Zhang, and Donglin Wang. DARA: Dynamics-aware reward augmentation in offline reinforcement learning. In ICLR , 2022.

Fan-Ming Luo, Shengyi Jiang, Yang Yu, Zongzhang Zhang, and Yi-Feng Zhang. Adapt to environment sudden changes by learning a context sensitive policy. In AAAI , 2022.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR , abs/1707.06347, 2017.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In ICML , 2018.

Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In ICML , 2020.

Xi Chen, Ali Ghadirzadeh, Tianhe Yu, Yuan Gao, Jianhao Wang, Wenzhe Li, Bin Liang, Chelsea Finn, and Chongjie Zhang. Latent-variable advantage-weighted policy optimization for offline RL. CoRR , abs/2203.08949, 2022.

Wenxuan Zhou, Lerrel Pinto, and Abhinav Gupta. Environment probing interaction policies. In ICLR , 2019.

Xiong-Hui Chen, Yang Yu, Qingyang Li, Fan-Ming Luo, Zhiwei (Tony) Qin, Wenjie Shang, and Jieping Ye. Offline model-based adaptable policy learning. In NeurIPS , 2021.

Benjamin Eysenbach, Shreyas Chaudhari, Swapnil Asawa, Sergey Levine, and Ruslan Salakhutdinov. Off-dynamics reinforcement learning: Training for transfer with domain classifiers. In ICLR , 2021.

Haoyi Niu, Shubham Sharma, Yiwen Qiu, Ming Li, Guyue Zhou, Jianming Hu, and Xianyuan Zhan. When to trust your simulator: Dynamics-aware hybrid offline-and-online reinforcement learning. CoRR , abs/2206.13464, 2022.

Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel. Constrained policy optimization. In ICML , 2017.

Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. CoRR , abs/1805.00909, 2018.

Finale Doshi-Velez and George Dimitri Konidaris. Hidden parameter markov decision processes: A semiparametric regression approach for discovering latent task parametrizations. In IJCAI , 2016.

Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. J. Artif. Intell. Res. , 76:201-264, 2023.

Luisa M. Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. VariBAD: A very good method for bayes-adaptive deep RL via meta-learning. In ICLR , 2020.

Jiachen Yang, Brenden K. Petersen, Hongyuan Zha, and Daniel M. Faissol. Single episode policy transfer in reinforcement learning. In ICLR , 2020.

Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In ICML , 2017.

Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta-reinforcement learning with advantage weighting. In ICML , 2021.

Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In ICML , 2017.

Jing-Cheng Pang, Tian Xu, Shengyi Jiang, Yu-Ren Liu, and Yang Yu. Sparsity prior regularized q-learning for sparse action tasks. CoRR , abs/2105.08666, 2021.

Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In ICML , 2019.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In NeurIPS , 2020.

Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. CoRR , abs/1911.11361, 2019.

Riashat Islam, Komal K. Teru, and Deepak Sharma. Off-policy policy gradient algorithms by constraining the state distribution shift. CoRR , abs/1911.06970, 2019.

Qiang Liu, Lihong Li, and Ziyang Tang and ´ Dengyong Zhou. Breaking the curse of horizon: Infinite-horizon off-policy estimation. In NeurIPS , 2018.

Ofir Nachum, Yinlam Chow, Bo Dai, and Lihong Li. Dualdice: Behavior-agnostic estimation of discounted stationary distribution corrections. In NeurIPS , 2019.

Tanmay Gangwani and Jian Peng. State-only imitation with transition dynamics mismatch. In ICLR , 2020.

Fangchen Liu, Zhan Ling, Tongzhou Mu, and Hao Su. State alignment-based imitation learning. In ICLR , 2020.

Tianwei Ni, Harshit S. Sikchi, Yufei Wang, Tejus Gupta, Lisa Lee, and Ben Eysenbach. f-IRL: Inverse reinforcement learning via state marginal matching. In CoRL , 2020.

Shentao Yang, Yihao Feng, Shujian Zhang, and Mingyuan Zhou. Regularizing a model-based policy stationary distribution to stabilize offline reinforcement learning. In ICML , 2022.

Paul F. Christiano, Zain Shah, Igor Mordatch, Jonas Schneider, Trevor Blackwell, Joshua Tobin, Pieter Abbeel, and Wojciech Zaremba. Transfer from simulation to real world through learning deep inverse dynamics model. CoRR , abs/1610.03518, 2016.

Shengyi Jiang, Jing-Cheng Pang, and Yang Yu. Offline imitation learning with a misspecified simulator. In NeurIPS , 2020.

Emanuel Parzen. On estimation of a probability density function and mode. The annals of mathematical statistics , 33(3):1065-1076, 1962.

Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio. Generative adversarial networks. CoRR , abs/1406.2661, 2014.

Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-gan: Training generative neural samplers using variational divergence minimization. In NIPS , 2016.

Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y. Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: model-based offline policy optimization. In NeurIPS , 2020.

Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In IROS , 2012.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR , abs/2004.07219, 2020.

Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In NIPS , 2016.

Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforcement learning via probabilistic context variables. In ICML , 2019.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a context-aware method called SRPO (State Regularized Policy Optimization) to address the issue of training RL algorithms on data with dynamics shift, i.e., with different underlying environment dynamics. The authors reveal and exploit the property of similar stationary state distributions in environments with similar structures but different dynamics. The learned stationary state distribution in SRPO is used to regularize policies trained in new environments, leading to more efficient data reuse and cross-domain ability. The paper also provides a lower-bound performance guarantee on policies regularized by the optimal state distribution in other dynamics. The experimental results show that SRPO can be a flexible add-on design that makes context-based RL algorithms far more data efficient and significantly improve their overall performance.\n\nSTRENGTHS:\n+ The paper is well-structured and easy to follow, with relatively rigorous experimental validations and inspiring theoretical analyses.\n+ The proposed SRPO algorithm is a context-aware method that tackles the challenge of training RL algorithms on data exhibiting diverse dynamics shifts, a very general problem in many real-world scenarios.\n+ The authors have identified and effectively leveraged the observation that environments with similar structures but different dynamics possess similar optimal stationary state distributions. This discovery facilitates efficient data reuse and demonstrates improved performance in both online and offline settings, as evidenced by the experimental results.\n+ The paper makes a valuable theoretical contribution by providing a lower-bound performance guarantee on policies regularized by the optimal state distribution in alternative dynamics.\n\nWEAKNESSES:\n+ Please refer to the subsequent question section.\n\nQUESTIONS:\n##### Method:\n+ While the state regularized objective in SRPO is indeed well-motivated and illustrated by examples in the Inverted Pendulum environment, it would benefit from additional theoretical evidence to firmly establish the assertion that environments with similar structures but different dynamics inherently share similar optimal stationary state distributions. It should be noted that such a claim may not always hold true in some environments, and further theoretical investigation is expected to determine the generalizability of this finding.\n+ Though the sample-based surrogate mechanism for state density ratio estimation is relatively well-founded, there is still some confusion regarding the inability of $D_{\\text{fake}}$ with lower state values to capture the stationary state distribution $d_\\pi$.  It remains a challenge to handle $D_{\\text{fake}}$ effectively, as it tends to underestimate the quality of the training policy and may even introduce a mixture of $d_\\pi$ and $d_{\\pi_\\beta}$  induced by the behavior policy in static datasets in offline settings.\n##### Experimental results:\n+ DARC and DARA objective is derived from constrained optimization with dynamics regularization, instead of state regularization in SRPO, which both decouple the regularization term with policy that might vary in environments with dynamics discrepancies. It would be beneficial and interesting to provide more extensive analyses or explanations about these two types of regularization and demonstrate in what situations SRPO can perform larger advantage and why SRPO can work with more cases. It might also accout for the results of offline experiments that \"MAPLE+DARA\" offers very competing performance against \"MAPLE+SRPO\" on halfcheetah and walker tasks.\n+ Further elaboration is needed to support the statement that SRPO enables efficient data reuse, specifically regarding reduced data consumption from source and target domains.\n+ If there are any specific design considerations aimed at stabilizing the training process, it would be beneficial to include a detailed list of implementation details related to discriminator training.\n\n##### Typos:\n- line 139 in Appendix: repeated \"static\".\n\nLIMITATIONS:\nThe authors have adequately and precisely addressed the limitations in this work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-17"}

comments: SUMMARY:
The paper considers leveraging existing data in RL but with dynamics shifts. The proposed method augments the reward with a density ratio estimation term that encourage the policy to recover the optimal state distribution under the distribution shift. The theory motivates the algorithm design by showing that under certain assumptions, optimal policies under similar dynamics have close stationary state distribution. Finally, the experiments show that the proposed method indeed works well in practice in both online and offline setting with dynamics shift.

STRENGTHS:
1. The paper is well-written and easy to follow.

2. The proposed method can be easily added to some previous methods.

3. The paper provides motivating theory results to back up the design of the practical algorithm.

4. The design of the density ratio estimation design is well motivated in Section 3.3.

5. The paper provides good coverage of the experimental setting (both online and offline and includes several environments in the mujoco benchmark). The result is pretty good compared with previous baselines.

WEAKNESSES:
1. The theory section (section 4) only provides some motivation for the algorithm design, instead of actually analyzing the algorithm.

2. The Lipchitzness of the reward function and dynamics still seems pretty strong even with section A.5: it should argument not only the Lipchitzness condition holds, and also the coefficients should be reasonably small.

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents an approach to RL/policy learning in the case where the dynamics parameters are potentially time-varying. The approach relies on the view that the optimal state distribution is similar across (small) variations in the underlying parameters, subject to assumptions. Thus, the authors propose an approach that constrains the distance between the realized state distribution under the trained policy and that of the optimal policy. Because the optimal policy state distribution is unknown, the authors propose to approximate the ratio of state distributions via discriminative methods, which in turn rely on an RL-as-inference formulation that leverages the value function.\n\nSTRENGTHS:\n- The problem is reasonable, and guaranteeing (in some sense) reasonable RL performance with time-varying parameters is an important step toward practical robustness.\n- The method is reasonable and original.\n- The experimental evaluation is reasonably comprehensive.\n\nWEAKNESSES:\n- The underlying motivation of similar state distributions is questionable. The similarity of these distributions is said to arise when the dynamics are similar (in this case, as reflected by (T, eps) notation) and when the reward functions and dynamics functions are Lipschitz. Then, the reward difference is approximately bounded by the products of the Lipschitz constants and the eps difference between dynamics functions. The authors claim these assumptions are mild---on the contrary, these conditions are quite strong and the resulting bounds are weak. In general, the idea of targeting the same state distribution doesn't seem reasonable for many real world systems, especially since a common use-case of RL is to handle non-smooth dynamics. \n- The approach is quite complex. The addition of a value-dependent discriminator is a substantial addition to standard RL workflows, and this added complexity will likely be a barrier to adoption. The authors should consider investigating simplified approaches---perhaps motivated by the same goals---that are more likely to see adoption. \n- Most importantly, the difference MAPLE + SRPO over MAPLE is extremely minor. Indeed, across many tasks the difference in performance is not statistically significant. In general, the combination of the high added complexity and the small resulting performance difference makes it difficult to recommend acceptance of the paper.\n\nQUESTIONS:\n- Is it possible to avoid the training of new discriminator while still obtaining the benefits of state distribution regularization?\n- Is it possible to generalize analysis beyond smooth systems?\n\nLIMITATIONS:\nThe conditions stated in the paper (discussed above) are described as mild. The conditions are not mild, and the limitations of these assumptions should be discussed in more detail.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper aims to improve sample efficiency for RL algorithms trained in environments with dynamics shift. The core motivation is that, for different environments and their corresponding optimal policies, the state visitation distribution could be similar. Building upon this motivation, the authors propose a constrained policy optimization (CPO) that encourages the learned policy to have a similar stationary state distribution given by the optimal state distribution. For theoretical analysis, the authors propose homomorphous MDPs, which require two MDPs to reach the next state with the same support given the same current state and action. The result is that, given the reward and dynamics functions of MDPs having Lipschitz properties, if a policy $\\hat{\\pi}$ has a similar stationary state distribution with the optimal policy in one MDP $M$, $\\hat{\\pi}$ will have a lower-bound performance guarantee in all MDPs that are homomorphous with the MDP $M$. Experimental results show that SRPO outperforms existing state-of-the-art methods in both online and offline RL settings.\n\nSTRENGTHS:\nThe motivation that, for similar MDPs, their optimal policies should have similar stationary state distributions is very intuitive and makes a lot of sense. This motivation is well supported by both theoretical analysis and experimental results.\n\nWEAKNESSES:\nPlease see the questions section.\n\nQUESTIONS:\n1. In equation 5, we use $R_\\max$ directly. What if $R_\\max$ is unknown? I assume we can estimate $R_\\max$ simply by running several rollouts, or set $R_\\max$ manually from domain knowledge. And what if the maximum reward we can obtain is different (or even has a different scale) for different states? E.g., we have $R_\\max$ as 100, but for state $s_1$ the maximum reward is only 0, and for state $s_2$ the maximum reward is 100. Given equation 5, we always subtract $r(s_t, a_t)$ by 100 for both $s_1$ and $s_2$, but this seems unfair for $s_t$.\n2. Line 159, there is no D.3 in the appendix.\n3. In some cases, the assumption that the optimal policy for similar MDPs has similar stationary state distributions might not hold. For example , driving a car (with a fixed power level) on roads with different slopes, and the reward function is the velocity of the car. Apparently, the maximum velocity on different roads can be different, thus affecting the stationary state distribution. We can find more examples like this, only if the stationary state distribution for the optimal policy is strongly correlated with the MDP parameter. We do not expect an algorithm to be applicable in all cases, but I think we should clearly state the scope of this algorithm.\n4. In both Figure 1 and the experiment section, we use a discrete set of MDPs. Several papers mentioned in line 104 \"MDPs with Different Dynamics\" indeed consider the case of continuous sets of MDPs, like VariBAD. For VariBAD, the generalization between different MDPs is carried out by the continuity of policy w.r.t. MDP dynamics. And it seems inappropriate to conclude \"data with different dynamics are used in an ad hoc manner. The large amount of data with different dynamics are discarded during training, leading to poor sample efficiency (line 38)\". Data collected from one MDP is helpful for other MDPs that are close to it. Can we 1) run experiments on continuous sets of MDPs, 2) make the discussion in line 37 more accurate? And if time permits, consider the VariBAD baseline.\n\nLIMITATIONS:\nThe authors did not address limitations or potential negative societal impact of their work. No further improvement should be made.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}