State Regularized Policy Optimization on Data with Dynamics Shift
Authors: Zhenghai Xue
Qingpeng Cai
Shuchang Liu
Dong Zheng
Peng Jiang
Kun Gai
Bo An
Venue: NeurIPS 2023 poster (2023)
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
Zhenghai Xue!
Qingpeng Cai Shuchang Liu?
Dong Zheng?
Peng Jiang" Kun Gai Bo An
Nanyang Technology University, Singapore
2 Kuaishou Technology Unnffliated
zhenghai001@e ntu. edu.sg boanêntu. edu. sg gai kwnkqu com com
fcaiqingpeng, wwwshuchang.,... Lushuuchaggg zhenggong, jianggennnnaaaa...... @kuaishuu..ooocomomm com
Abstract
In many real-world scenarios, Reinforcement Learning (RL) algorithms are trained on data with dynamics shift, i.c., 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 data are used ad hoc, and policies trained for one dynamics cannot benefit from data collected in all other environments with different dynamics. 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 new environment, leading to the SRPO (State Regularized Policy Optimization) algorithm. To conduct theoretical analyses, the intuition similar environment structures is characterized by the notion homomorphous MDPs. We then demonstrate a lower-bound bound performance guarantee on policies regularized by the stationary state distribution. In practice, SRPO can be add-on module context-based algorithms in both online and offline RL settings. Experimental results show that SRPO can make several context-base algorithms far more data efficient and significantly improve their overall performance.
1 Introduction
Reinforcement Learning (RL) has achieved great success in solving challenging sequential decision- making problems 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 varying and heterogeneous user preferences 3 4. A robot arm may operate in different scenarios with different joint frictions and medium densities [.. these cases, the agent has to work with the trajectory data from different environment dynamics, i.c., 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 g can easily be misled by different environment dynamics and fail to train 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 to train context encoder 19] [0 [ to associate the environment dynamics with a latent variable. The policy then trained with the latent variable as an additional input [12). One issue with this
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

practice is that policies conditioned on a specific latent variable can only learn from data collected in the environ- ment corresponding to that latent variable. In other words, data with different dynamies are used in an ad hoc man- ner. The generalizability of context encoders relies the expressive power neural networks. However, neural networks are prone to overfit and behave poorly when ex-
Performance in Ant Environment
x10²
Algorithm
CaDM
Return
CaDM+SRPO
trapolating. As an example, we benchmarked CaDM [9].
which is one of the context-based algorithms, under Ant
environments with different gravities and display the re- sults in Fig. Although it can outperform PPO [7] due
Number of Different Dynamics
to its adaptability from context encoders, CaDM fails to constantly improve its performance with more data from different environment dynamics. To mitigate the prob- lem of inefficient data use, there leverage Importance Sampling (IS) [3 S [4]. 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 the target en- vironment and vice versa. Compared to training context
Figure 1: Performance comparison of PPO 77.. CaDM and CaDM+SRPO in the Ant environment, where SRPO is our proposed state regularized policy op- timization method. Details of the experi- ment setup are in Sec. 5.1
attempts that
encoders. IS-based methods manage 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. 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 set of environments with similar structures and different environment dynamics. For example, given heterogeneous preferences of users, 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 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 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 unified and efficient approach to exploiting these data.
Summarizing these ideas, we propose the SRPO (State Regularized Policy Optimization) algorithm. SRPO works as an add-on module in both online and offline context- based RL algorithms such as CaDM (9] and MAPLE [12] 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-ard context-based RL algorithms. We also perform ablation studies to demonstrate the effectiveness of each component the SRPO algorithm.
2 Backgroud
2.1 Preliminaries
Markov Decision Process (MDP) can be defined by tuple (S. A.A.ATTr,yy Po)), where is the state space, A is the bounded action space with actions a E (-1, [(s' E [0, and r(s, s') ) E [- --nax, Rmax] are the transition and reward functions. 7 E (0, 1) is the discount factor and po(s) is the initial state distribution. In MDPs with deterministic transitions, we denote T(x, a) as the transition function with slight abuse of notation, and (T, as (T" [T(s, a) -' (s, <, s, Vs E S, E A) which is the e neighbourhood of T.. RL aims maximizing the accumulated return of policy

---

(x)= Ex, Σ y /r(s,, ai) where the expectation is computed with so Pa. at (-|(s)).
and 81+1 T())tt;))) The optimal policy n² is defined कạ. arg max (r)). In MDP
with a policy т, the Q-value (7 a) denotes the expected return after taking action a at state Σ 70r(8,, as) = 0, = = The value function is defined as V*()) =
8: (7(;,0)= F
1==
B-->(11s) (11s) (7 a) with (s) being the shorthand for ViJT satisfies the optimal Bellman
Equation (s) = max r(s, a) + ".BBerr(l Vi'(s). We can also define the stationary state
distribution (also known as state occupation function) as (s) == (1 - 7) E' (st = with (s) being the shorthand for di (s).
The Hidden Parameter Markov Decision Process (HiP-MDP) captures a class of MDPs with dif- ferent transition functions and the same reward function by introducing set of hidden parameters. Specifically, an HiP-MDP is defined by a tuple (S, A, θ T, 7; po), where θ is the space of hidden parameters. The transition function (a(s'|s (s, a, 0) is parameterized not only by states and actions, but also by a hidden parameter 0 sampled from O The action gap of an HiP-MDP is defined as A min min min Vi, (s) Q7, (s,a), a). which reflects the minimum gap between an optimal AE sES (s)
action and all other sub- sub-optimal actions.
2.2 Related Work
MDPs with Different Dynamics The setting of HiP-MDP 17 was proposed model 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 Unfortunately, these methods train dynamics- specific 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 [33 14 S according to the transition probabilities in the test environments, which is often not available in real-world scenarios. Finally, meta-RL. RI. algorithms 21 22 can adapt to environments with new dynamics through fine-tuning on a small amount data from the test environment. In contrast, our method can be directly applied to new environment dynamics by making zero-shot transfer.
the transition function [9] or policy network [20] 12
Behavior Regularized Methods
The idea of constrained policy optimization (CPO) [S] 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:
max EaDD Ea (-19) Q(+,a'))
s.t.
EEED [D(DDD((TTTT [(T((1s)) <E.
(1)
where D is the replay buffer, â is the regularizing policy, and Ó is certain distance measure. Maximum-Entroyy RL 23 can be considered as CPO with a uniform policy distribution. The sparse action tasks (24 can be solved by CPO with 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 's policy distribution. However, data sampled from environments with different dynamics can have distinct optimal policies, as illustrated in Sec. 3. In such cases, â in Eq. (T) may include policies that do not match the current environment, and can therefore be misleading. So behavior regularization in Eq. (I) 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 di of policy ı and dynamics 7' is an important feature that can measure the differences in policies and transition functions. It has already been exploited many researches. Off-Policy RI., Islam et. al [28] estimates the stationary state distributions of both the current policy and the mixed buffer policy. It
3

---

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 (II.),
state-only II. 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 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 paper is similar to some researches sim-to-real 35 36]. They propose match the next state distribution in the imperfect simulator and the real environment with inverse dynamics model. They implicitly relies 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.
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 See. 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 op- timal state distribution is similar across envi- ronments. Consider example of the Inverted Pendulum environment in Fig. 2 We train two policies in the environments with gravities of 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 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 pen- dulum angle, while the action distribution has different density peaks. It demonstrates that state distribution of data generated by the optimal policy can be similar regardless of the environment dynamics, and therefore can serve reference distribution to regularize the train- ing policy in environments with new dynamics. More demonstrating examples can be found at
4.0 State Density with Gravity
State Donsity with dravity 10
4.8
0.0
1.6
4.0
2.4
25
Angle
Angle
Torque
Tarque
Figure 2: Visualization of state and action densities in data sampled from the Inverted Pendulum envi- romment with gravity 5 and 10. Under both grav- ities, the state distribution has high density with low pendulum speed and small pendulum angle. Meanwhile, the action distribution has different peaks in density under different gravities.
Appendix B.3
3.2 State Regularized Policy Optimization
Based the intuition of informative optimal state distribution, we develop novel technique that regulates RI. algorithms to generate a stationary state distribution that close to the optimal one. Specifically, we propose the following constrained policy optimization formulation:
8
max Eg Egyay Tu
Σ 1=
rrr (8t,
s.t. DKL (((()))))) < E,
(2)

---

where {(s) is the optimal state distribution in other environment dynamics. By introducing the stationary state distribution, the optimization problem defined Eq. extends the regularization of in-distribution data to data with distribution shif
We solve Eg. by casting it to the following unconstrained optimization problem via Lagrange multipliers
Σ 7 (r(s), az) + log <(G))
d,
where À > 0 is the Lagrangian Multiplier. It is noteworthy that in addition to the multiplier term, the only difference of Eq. (3) and the geward-maximizen
objective of RL is that the logarithm of <(ss) is added to the reward term r(*t, at). Therefore, one can easily probability density ratio À log (=)) 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 [, because obtaining the optimal state distribution during online training given suboptimal offline dataset is infeasible. Also,
((s)
d, (s) is intractable if the state space is continuous. Motivated by recent advances in adversarial
training 38 39). we propose sample-basee based surrogate for the density ratio ((s)
d.(*)
Proposition 3.1. In a GAN, when the real data distribution is ((s) and the generated data distribution (s), the output of the discriminator D(s) follows
D(s) ((s)
1 ((s) da (s)
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 the optimal state distribution
@1
02
03
(s) and away from (s), which is sort of the average state distribution. Motivated by [16), we model state optimality by a variable O, As shown in Fig. ] we regard the state s; MDP as a hidden state in a Hidden Markov Model (HMM), and introduce the binary observation state O,. O = 1 denotes that is the optimal state at timestep t. The observation model is given by
91
82
83
O
Figure 3: HMM in MDP with optimality variables O,.
p(O,(s:) max
(r(0))/(((ssa))) Rmax)]]
(5)
We can therefore compute the state density ratio ((x)
[+)
((s) dx (s) (s)
p(O (s, x ) (s)
Eafp(Oua11ls: 1 T)p(O, m)
(6)
(s)
P(O0:o0177) (d)
p( (T)
where the second equation follows the Bayes' law. The last term is related to the forward probability Ore (xx) = p( P(OO.t-1l 18t, 100t-1185..
T)) and backward probability 3r (sx) = p(Orsco T) the HMM. We discuss in Appendix A.2 that (st) is positively related to soft version of MDP's state value V (st). Also, (st)) 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 higher state value ((s) than average. With this idea, we are able to build training samples for the discriminator ((*).
similar form of Eq. (See Eq. in Sec. 2.2 employed in Offline RL algorithms ensure conservative policy updates. But restricts the training data to he sampled from the same environment. The detailed derivations of the Lagrangian can be found in Appendix A.
5

---

Algorithm The workflow of SRPO on top of MAPLE 12
1: Input: 0pp a context encoder parameterized by yi Adaptable policy network To parameterized by 0; Adaptable value network Ve parameterized by v;; Offlin dataset Don; Rollout horizon H; State partition ratio p; State discriminator D5 parameterized by 6; Regularization coefficient X.
2: for 1, 3........ do
3: for ( 1....... H. do
4: Sample ž, from Uxp (= at-10-1110 1t- and then sample ar from my (a 8t, zx).
5: Rollout and get transition data ($t41, Tty dr+1, St, oes ze). Then add it to Pooloou
6: end for
7:
Update the context encoder 6.p according MAPLE.
8:
Sample a batch busch from Doir and Pollou and rank them by their state-values estimated by
Ve: Add p|Dbuchh states with higher state-values Dreat and the others Diuke
9:
Train the discriminator D5 with nll loss.
10: For one-step transition (S++1, Tt; d,+1, 8t, dt, ze) in Dratch update r, with r + A Da(ss)
11: Use the updated Daach and SAC to update the policy and value network parameters @ and tb.
12:
end for
3.4 Practical Algorithm
Summarizing the previous derivations, we obtain practical reward regularization algorithm, termed as SRPO (State Regularized Policy Optimization) to leverage data with dynamics shift. We select the Offline RL, as the
MAPLE 12 algorithm, which is one of the SOTA algorithms in context- base algorithm. The detailed procedure of MAPLE+SRPO is shown in Alg. I After preparing the dataset in a model-based Offline RL style 40 12], we sample a batch data from the dataset, obtain portion of states with higher rewards and add them to the dataset Dreat Diake is similarly generated by states with lower rewards (line 10). We set p = 0.5 offline experiments with medium-expent level of data. p 0.2 is set in all other experiments. Then classifier discriminating data from (())
the two datasets is trained (line 11). It estimates the logarithm of the state density ratio À log d, (s)* which is added to the reward rt (line 12). A is regarded as a hyperparameter with values 0.1 or 0.3. The effect of À is investigated in Sec. S.c.5.3 The procedure of the online algorithm CaDM [S+RR similar to MAPLE+SRPO, where the datasets Dreal and Diske 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 Sec. 3 The notations are introduced Sec. Sec. [2]] and proofs can be found in Appendix A.4 We first show in Thm. 4.2 that the performance of policy can be lower-bounded when its stationary state distribution is close to 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,0,T, Po), consider hidden T(s'ls, a, 0,), s") E S x A 8, = 1,2. If + Σ T2(s'ss,, s, 0 for all s, E S. MDPs (8, T. and
parameters 01, O2 E O. Let (s's
Σ Ti (s'ls, a T.,. are referred to homomorphous MDPs
0
aEA
In this definition,
Σ T(s'ls, > 0 means state s' can be reached from so the equivalence of aEAA
non-zero transition probabilities refers to the same reachability from to s'. Such condition holds in 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 Å -Lisschiiz and the dynamics funetion T(s, a) w.r.t. the action a is Xg-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 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 theorem as follows.
Theorem 4.2, Consider two homomorphous MDPs with dynamics T and T', IfT E (T. all
learning policy āt such that DRL (d7 (-)l77)))))))))))) & Es. we have
Am 221 + /2
(x( > Tr (T7)
1 1-y
The theorem implies that if policy ø has 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 z is related to the dynamics shift mm. the LL-eivereeece of two stationary state distributions Es, and the effective planning horizon T- 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 pre-defined RI. task, the discrepancy between stationary state distributions E, is determined by the learning policy á and can be optimized during training to obtain a better performance lower-bound. We also discuss Appendix A.5 how tight Eq. (7) is in terms of the effective planning horizon I compared with some similar performance bounds.
With an additional assumption on the action gap A (defined in Sec. 2.I.. 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', IfI E (T, and the action gap A follows A > 2--))a)>25 for all S we have di) = dz. (s).
1-7
The assumption is mild and holds in many scenarios. For example, 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 A. recommendation tasks we are hardly concerned with what items we recommend (the action), long the recommendation outcome (the state), i.e., the users' experiences are good enough, leading to a small A1. The condition of large enough action gap holds in these situations.
Experiments
In this section, 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 [1] and generate environments with different transition functions. We train the CRDM+SRPO and the MAPLE+SRPO algorithm proposed in Sec. 3.4.and make comparative analysis with baseline algorithms. To answer question (2), do ablation studies 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), 10.7, and 1,..00.,,1.0,1.33.11 in experiments with 1, and 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 proportion to the number of environment parameters. Therefore, the agent has access to the same amount of training data
7

---

FaDK-SAPo HalfCheetah Different smics!
Pendulum 13 Different Dynamics)
Art Different Dynamics)
Humanoid Different Dynamics)
-0.4
X20
X10
X18
Ki0
3.0
-9.6
2.5
2.0
2.0
0.8
1.5
1.5
-1.0
1.0
0.4
0.5
0.5
1.4
1.5
0.0
Pendulum Dynamics)
Dynamics)
HallCheettah Dyn amnics)
Humanoid Dynamics)
K20*
1.2
1.0
-0.8
-100
0.4
0.2
Timesteps(M)
Timesteps(M)
Timesteps(M)
Figure Results of online experiments on MuJoCo tasks. The comparison is made between CaDM+SRPO and baseline algorithms PPO, CaDM. Our CaDM+SRPO algorithm has the best overall performance experiments with and different environment dynamics. The curves show the average return on 4 random seeds and the shadow areas reflect the standard deviation.
Table 1: Results of offline experiments on MuJoCo tasks. Numbers are the normalized scores according to the D4RL paper (42). ME, M, MR and R correspond to the medium-expert, expert, medium-replay and random dataset, respectively. The evaluation is done on policies at the last iteration of training. averaged over four random seeds. The number after + is the standard deviation. Our proposed MAPLE+SRPO algorithm has the best performance in 8 of 12 tasks and the highest overall performance.
CQL (Single Env)
GAIL
CQL
MOPO
MAPLE
MAPLE +DARA
MAPLE +SRPO(Ours)
Walker2d-ME
1.11
0.21
1.03-1:011 10
0.25=10.18
0.55=0.21
0,80-10.02
0.66-1:0.08
Walker2d-M
0.79
0.15++0.06
0.78++0.01
0.23++0.34
0.82=0.01
0.83+0003
0.84::0.05
Walker20-MI
0.27
0.00:10.00
0.07:1:0:00
0.00:1:000
0.16==0.22
0.17:20.01
0.17::0.02
Walker2d-R
0.07
0.00.10.00
0.03:10.00
1.00
0.22 0.00
0.22.10.00
0.22.1 0.00
Hopper-ME
0.98
0.04
0.32:t:0.14
0.00
0.96=0.14
0.96::0006
0.98
Hopper M
0.58
0.00:1:0.00
0.57:10.116
0.01 0.000
0.78
0.40:1:0.05
1.03
Happer-MR
0.46
0.00-1:..0
0.14:1:0.02
0.01
0.91-
1.02
Hopper-R
0.11
0.00+t:000
0.11=:0000
0.01
0.13=0.00
0. 03:+0.001
0.32-+0.02
HallChhstah-ME
0.62
0.36+F0.06
0.03-4:0044
-0.03-1:0.00
0.50=0.06
0,50=t0.00
0.63+00.01
HallCheetab-M M
0.44
0.25:1:002
0.43:1:0.03
0.38:10,28
0.62==0.01
0.67:10.03
0.63:1:0.01
HalfCheetah MR
0.46
0.18.10.11
0.46:10.00
-0.03
0.52.10.00
0.53
0.50
HalfChestah--
0.35
0.14:00022
0.01
-0.03-10.00
0.22=00.33
0.21
0.24:1:.01
Averago
0.52
0.11
0.33
0.068
0.53
0.61
certain value of simulator gravity. We also consider the shift of medium density and body mass in offline experiments show SRPO's robustness different forms of dynamies shift.
To perform comparative analysis, we choose CaDM 9 and PPO 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, choose MOPO [40), MAPLE [12] and CQL (26) 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 Fig. 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 dynamies shift. Its performance with 5 different dynamics is lower than that with 3 dynamics. In contrast, our proposed SRPO

---

Table 2: Results of ablation studies in offline experiments.
Table 3: Performance comparison of differently regularized policies.
MAPLE +SRPO
Behavior Regularizing
Random Partition
Fixed
Original Hopper Env
10x
Density
Walker2d
0.47
0.45
0.39
0.40
Hopper
0.83
0.68
0.56
0.79
Random
121.3
44.15
HalfCheetah
0,51
0,50
0.45
0.40
Medium
2178
913.3
Average 0.61
0.54
0.47
0.53
Expert 3819
3748
algorithm leads to better performance on top of CaDM in accordance with more training data. significantly outperforms the original CaDM algorithm in environments with 5 different dynamics. The performance comparison the Pendulum environment also in accordance with the motivating example in Sec. 3.1 More results of online experiments shown in Appendix B.2
Offline Experiments The results of offline experiments shown in Tab. The column of "CQL Single" refers to the evaluation score in the CQL (26 paper, where the policy is with data from single static environment. Without the mechanism of context-based encoders, GAIL 43 CQL and MOPO 400 cannot handle data with distribution shift and show a performance drop. MAPLE [12 and MAPLE+DARA [5] only achieve marginal performance improvement with respect 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 P algorithm also has 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 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 lower evaluation score on average than policies trained with the original SRPO algorithm in all environments. This because environments with different dynamics do 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 classifier discriminate states with higher values from lower values. We also train another classifier discriminating random binary partition of states. Ablation results show huge performance drop, which verifies the effectiveness of the classification-hate surrogate mechanism. We also evaluate MAPLE+SREE with fixed value of the hyperparameter
A. A 0.1 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 Discriminators We first train a discrimi- nator D3 according to Alg. [ Then set of states is sampled from the D4RL [42 dataset and classified into two sets ac- cording the output of Ds. The average values of the two sets states are compared Fig- 5 As shown in the figure,
State Values With Different Output
1.2
High
Low
Normalized Values
2.0
states with higher D, outputs also have higher values in all three environments. It means that the trained discriminator Ds successfully identify states with high values from those with low values. Therefore, its output can be good
0.8
0.6
surrogate for the density ratio {(e) in Sec. 3
(e)
Hopper HaifCheetah Envrironment
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
Figure 5: Comparison of values on states with high and low output of the discriminator D.

---

used to estimate the stationary state distributions, which regularize the learning policies in new environment with different dynamics. The results are shown Tab. where the expert policy 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 regularizer. After discussions 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 iynamics-invarian
sample- based surrogate, we propose the SRPO algorithm which can be
dynamics invariant state distribution, as well as the constrained policy optimization formulation.
Limitations and Future work The theoretical analyses 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 - Pre- positioning (IAF-PP) Funding Initiative and Ministry of Education, Singapore, under its Academic Research Fund Tier (RG13/22), Any opinions, findings and conclusions recommendations expressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Singapore.
References
[1] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin A. Riedmiller. Playing atari with deep reinforcement learning. CoRR, abs/1312.5602, 2013.
[2] David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Vedavyas Panneershelvam, Mare Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy P. Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mas- tering the game of go with deep neural networks and tree search. Nat., 529(7587):484488) 2016.
[3] 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.
[4] 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.
[5] Jinxin Liu, Hongyin Zhang, and Donglin Wang. DARA: Dynamics-aware reward augmentation in offline reinforcement learning. In ICLR, 2022.
[6] 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.
[7] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, 02s/1707.006344 2017.
[8] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critie: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In ICML, 2018.
10

---

[9] Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In JCML, 2020.
[10] Xi Chen, Ali Ghadirzadeh, Tianhe Yu, Yuan Gao, Jianhao Wang, Wenzhe Li, Bin Liang, Chelsea Finn, and Chongjie Zhang. Latent-variable advantage-weight policy optimization for offline RL. CoRR, abs/2203. 308999, 2022.
[11] Wenxuan Zhou, Lerrel Pinto, and Abhinav Gupta. Environment probing interaction policies. In ICLR, 2019.
[12] Xiong-Hul Chen, Yang Yu, Qingyang Li, Fan-Ming Luo, Zhiwei (Tony) Qin, Wenjie Shang, and Jieping Ye. Offline model-based adaptable policy learning. In NeurlPS, 2021.
[13] Benjamin Eysenbach, Shreyas Chaudhari, Swapnil Asawa, Sergey Levine, and Ruslan Salakhut- dinov. Off-dynamics reinforcement learning: Training for transfer with domain classifiers. In ICLR, 2021.
[14] Haoyi Niu, Shubham Sharma, Yiwen Qiu, Ming Li, Guyue Zhou, Jianming Hu, and Xianyuan Zhan. When trust your simulator: Dynamics-aware hybrid anffine-and-online reinforcement learning. CoRR, abs/2206.1 13464, 2022.
[15] Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel. Constrained policy optimization. In ICML,, 2017.
[16] Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. CoRR, abs//805,,0 0005,0 2018
[17] Finale Doshi-Velez and George Dimitri Konidaris. Hidden parameter markov decision processes: A semiparametric regression approach for discovering latent task parametrizations, In IJCAI,
2016.
[18] 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-280 -264, 2023.
[19] 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 leamnng. In ICLR, 2020.
[20] Jiachen Yang. Brenden K. Petersen, Hongyuan Zha, and Daniel M. Faissol. Single episode policy transfer in reinforcement learning. In ICLR, 2020.
[21] Chelsea Finn, Picter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adapta- tion of deep networks. In ICML, 2017.
[22] Eric Mitchell, Rafacl Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta-reinforcement learning with advantage weighting. ICML, 2021.
[23] Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep evergy-based policies. In ICML, 2017.
[24] Jing-Cheng Pang, Tian Xu, Shengyi Jiang, Yu-Ren Liu, and Yang Yu. Sparsity prior regularized q-leamnin for sparse action tasks. CoRR, 10t.2/105006666 2021.
[25] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In ICML, 2019,
[26] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In NeurIPS, 2020.
[27] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. CoRR, abs/1911.11331. 2019.
[28] Riashat Islam, Komal K. Teru, and Deepak Sharma. Off-policy policy gradient algorithms by constraining the state distribution shift. CoRR, abs/1911.06970, 2019.

---

[29] Qiang Liu, Lihong Li, and Ziyang Tang and Dengyong Zhou. Breaking the curse of horizon: llfinite-horizon off- off-po estimation. In NeurIPS, 2018.
[30] Ofir Nachum, Yinlam Chow, Bo Dai, and Lihong Li. Dualdice: Behavior-agnostic estimation of discounted stationary distribution corrections. In NeurIPS, 2019.
[31] Tanmay Gangwani and Jian Peng. State-only imitation with transition dynamics mismatch. In ICLR, 2020.
[32] Fangchen Liu, Zhan Ling, Tongzhou Mu, and Hao Su. State alignment-based imitation learning. In ICLR, 2020.
[33] 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.
[34] Shentao Yang, Yihao Feng, Shujian Zhang, and Mingyuan Zhou. Regularizing model-based policy stationary distribution to stabilize offline reinforcement learning. In ICML, 2022.
[35] Paul F. Christiano, Zain Shah, Igor Mordatch, Jonas Schneider, Trevor Blackwell, Joshua Tobin, Pieter Abbeel, and Wojeiech Zaremba. Transfer from simulation real world through learning deep inverse dynamics model. CoRR, abs/1661003318,, 2016.
[36] Shengyi Jiang, Jing-Cheng Pang, and Yang Yu. Offline imitation learning with misspecified simulator. In NeurlPS, 2020.
[37] Emanuel Parzen. On estimation of probability density function and mode. The annals of mathematical statistics, 336));1065-107, 1962.
[38] Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. CoRR, abs/1406.2661, 2014.
[39] Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-gan: Training generative neural samplers using variational divergence minimization. In NIPS, 2016.
[40] 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.
[41] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: physics engine for model-based control. In IROS, 2012.
[42] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR, abs/2004.07219, 2020.
[43] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In NIPS, 2016.
[44] Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforcemmntt learning via probabilistic context variables. In ICML, 2019.
12


Reviews:
[{"review_id": "iWZEPVbvcp", "reviewer": "Reviewer_mfzT", "summary": "The 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.", "questions": "1. 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.", "limitations": "The authors did not address limitations or potential negative societal impact of their work. No further improvement should be made.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The 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.", "weaknesses": "Please see the questions section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback! We provide discussions on your concerns as follows.\n\nQ: **In equation 5, we use Rmax directly. What if Rmax is unknown or different across states?**\n\nA: In Eq. 5, $R_\\max$ serve as a normalization factor to guarantee the probability $p(\\mathcal{O}\\_t|s\\_t)$ is positive. In [Sergey 2018](https://arxiv.org/abs/1805.00909), $R_\\max$ is simply set to 0 without loss of generality. Also, what we really care about is the optimality of a state trajectory. When estimating the probability of a trajectory being optimal, single-state optimal probabilities are multiplied together. Different $R_\\max$ across states do not pose an issue when considering a whole state trajectory.\n\nQ: **Line 159, there is no D.3 in the appendix.**\n\nA: Sorry for this blunder. D.3 in Line 159 refers to Appendix C.3. We will change it in the revision.\n\nQ: **In some cases, the assumption might not hold. Consider the car driving scenario.**\n\nA: Thanks for pointing it out. We will add the following discussion in the revision. In the theory part (Section 4) of the paper, we provide the condition under which using optimal state distributions as regularizers will be helpful for policy training (Thm. 4.2). So theoretically SRPO is expected to work effectively in tasks following such condition. The condition contains the homomorphous property of MDP and two Lipstchitz properties, where the Lipstchitz properties hold in a large variety of environments, as elucidated in the general response. \n\nIn the car driving scenario the intuition of similar optimal state distribution does not hold, in that the homomorphous property is violated. Such property requires the reachability from $s$ to $s'$ is the same in the source and target environments. In the example of car driving, the power levels are fixed. Therefore, cars on slopes that are steeper can not accelerate as fast as cars on other slopes, leading to different reachability. Nevertheless, in most RL tasks the homomorphous property holds in at least states that are frequently visited, as demonstrated by the superior performance of SRPO in both online and offline experiments.\n\nQ: **Can we 1) run experiments on continuous sets of MDPs, 2) make the discussion in line 37 more accurate, 3) consider the VariBAD baseline?**\n\nA: We will change the discussion in line 37 to: \"The generalization ability of context encoders relies on the expressive power of neural networks. However, neural networks are prone to overfit and behave poorly when doing extrapolation\".  Also note that SRPO can provide context encoders with a generalization ability **orthorgorical** to that of neural networks. Both generalization effect can work together. Therefore, SRPO will also work well in a continuous set of MDP. Indeed, the comparison of PEARL+MAPLE and PEARL in the appendix is made on a continous set of Walker2d environments, where PEARL+MAPLE has better performance.\n\nRegarding the VariBAD baseline, we conducted additional experiments in Pendulum, Ant, HalfCheetah and Humanoid environments, each with 5 different dynamics. The results are shown in Fig. 1 of the one-page pdf in the general response. **VariBAD exhibits a better performance compared with CaDM and PPO baselines, but can not outperform our CaDM+SRPO algorithm**. We would also like to mention that as a general add-on module, SRPO can also work together with VariBAD, resulting in the VariBAD+SRPO algorithm.\n"}, {"review_id": "1shyMhm4nJ", "reviewer": "Reviewer_PXH4", "summary": "This 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. ", "questions": "- 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?", "limitations": "The 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. ", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback! We provide discussions on your concerns as follows.\n\nQ: **The similarity of similar state distributions is said to arise when the dynamics are similar and when the reward functions and dynamics functions are Lipschitz. These conditions are quite strong and the resulting bounds are weak.**\n\nA: We would like to respectfully counter with the argument of strong conditions and weak bounds. We restate the main performance bound here:\n$$\n\\eta_T(\\hat{\\pi}) \\geqslant \\eta_T\\left(\\pi_T^*\\right)-\\frac{\\lambda_1 \\lambda_2 \\varepsilon_m+2 \\lambda_1+\\sqrt{2} R_{\\max } \\sqrt{\\varepsilon_s}}{1-\\gamma},\n$$\n\nwhere $\\lambda_1$ and $\\lambda_2$ are Lipstchitz coefficients, $\\varepsilon_s$ is the error of state regularization $D_{\\mathrm{KL}}\\left(d_T^{\\hat{\\pi}}(\\cdot) \\| d_{T^{\\prime}}^*(\\cdot)\\right)$, and $R_{\\max }$ is the maximum possible reward. As clarified in the general response, $\\lambda_1 \\lambda_2$ and $\\lambda_1$ are more than 20 times smaller than $R_{\\max }$, making the term with $R_{\\max }$ dominate the error bound. Additionally, in similar performance bounds such as Thm. 1 in TRPO and Thm. 4.1 in MBPO, $R_{\\max }$ is multiplied by squared planning horizon $\\frac{1}{(1-\\gamma)^2}$. Thanks to the assumptions on Lipstchitsz conditions, our bound is stronger than these bounds in that $R_{\\max }$ is multiplied by $\\frac{1}{1-\\gamma}$.\n\nQ: **The underlying motivation of similar state distributions is questionable. 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\nA: We would like to emphasize that the assumption was made on the **$\\lambda_2$-inverse** Lipstchitz property of the dynamics function, rather than the more commonly used $\\lambda$-Lipstchitz property. This crucial distinction ensures that our assumption does not overly constrain practical tasks to have smooth dynamics, and the underlying motivation remains valid even in tasks with non-smooth dynamics.\nTo illustrate this point, consider the standard cliff walking task as an example. In this scenario, the optimal state trajectory involves the agent walking alongside the cliff towards the goal position. Even with a small change in action, the agent may fall off the cliff, resulting in a significant variation in the next state. Hence, this task exhibits non-smooth dynamics. Despite the non-smooth dynamics, the optimal state distribution remains consistent across different dynamics. It will always have high-density areas near the cliff.\n\nQ: **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.**\n\nA: Introducing an additional discriminator is a common and widely adopted practice in RL community, as demonstrated in various studies such as [LFIW](https://arxiv.org/abs/2006.13169), [ReMERT](https://arxiv.org/abs/2105.07253), [DARC](https://openreview.net/forum?id=eqBwg3AcIAK), and [DARA](https://arxiv.org/abs/2203.06662). In our paper, we incorporate the discriminator module into three different RL algorithms: CaDM, MAPLE and PEARL (in appendix). Remarkably, all three algorithms exhibit excellent performance, further highlighting the versatility and adaptability of the additional discriminator across different areas of RL. With respect to time complexity, we show the running time comparison of MAPLE, MAPLE+DARA (the baseline method), and MAPLE+SRPO (our method) in the following table. SRPO only needs about 15% of the extra training time and is more efficient than DARA.\n\n|  | Training Time (h) |  |  |\n| :---| :---: | :---: | :---:|\n|  | MAPLE | MAPLE+DARA | MAPLE+SRPO |\n| Hopper-v2 | 5.47 | 6.32 | 6.14 |\n| Walker2d-v2 | 7.20 | 8.48 | 8.25 |\n| HalfCheetah-v2 | 7.05 | 8.32 | 8.19 |\n\nIt is also crucial to note that the training procedure of the discriminator in our paper is fundamentally different from that in generative adversarial networks (GANs). Unlike GANs, which are notoriously challenging to stabilize during training, our approach does not involve a generator module. Consequently, the discriminator in our study serves as a simple classifier, leading to quick convergence and enhanced training stability, as shown in Fig. 2 of the one-page pdf. We will add these discussions and statistics in the revision.\n\nQ: **The difference of MAPLE + SRPO over MAPLE is extremely minor.**\n\nA: MAPLE+SRPO and MAPLE are algorithms designed for the domain of offline RL in environments with dynamics shift. In this setting, performance improvements from different algorithms are generally limited due to the data-collecting policy (or behavior policy) being directly trained in the target environments and thus not affected by dynamics shift.\n\n As shown in the following table from [H2O](https://arxiv.org/abs/2206.13464), [DARC](https://openreview.net/forum?id=eqBwg3AcIAK) achieved ~10% improvement over the baseline algorithm CQL. On top of DARC, H2O achieved ~11.5% performance improvement. Considering our approach, the MAPLE+SRPO algorithm demonstrates a 15% increase in performance compared to MAPLE and a 13% improvement over MAPLE+DARA. These results indicate a significant and reasonable enhancement achieved by the MAPLE+SRPO algorithm in the context of offline RL with dynamics shift.\n|  | CQL | DARC | H2O |\n| :--- | :--- | :--- | :--- |\n|Average Performance  | 5196 | 5715 | 6373 |\n| Improvement over baseline  | N/A | 9.98\\% | 11.5\\% |\n\nOn the other hand, in online RL the data-collecting policy is the same as the policy that is being trained.  When a policy cannot adapt to environments with dynamics shifts, it will struggle to collect high-quality data to train and have a very poor performance. The second row of Fig. 4 in our paper demonstrates the large performance gap of CaDM+SRPO with other baseline algorithms in online experiments.\n\n"}, {"review_id": "wqbOEq0Wtq", "reviewer": "Reviewer_ZrWb", "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.\n\n", "questions": "N/A", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and easy to follow.\n\n2. The proposed method can be easily added to some previous methods.\n\n3. The paper provides motivating theory results to back up the design of the practical algorithm.\n\n4. The design of the density ratio estimation design is well motivated in Section 3.3.\n\n5. 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.\n\n2. 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.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback! We provide discussions on your concerns as follows.\n\nQ: **The theory section (section 4) only provides some motivation for the algorithm design, instead of actually analyzing the algorithm.**\n\nA: The SRPO algorithm works by augmenting the reward function. Such augmentation changes the original MDP and it is usually hard to analyze the sample complexity or convergence rate of the new algorithm. These analyses are also absent in similar reward augmentation algorithms, such as \b[DARA](https://openreview.net/forum?id=9SDQB3b68K), [DARC](https://openreview.net/forum?id=eqBwg3AcIAK) and [H2O](https://arxiv.org/abs/2206.13464). Meanwhile, the experiments results can serve as empricial analyses of the SRPO algorithm.\n\nQ: **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.**\n\nA: Thanks for pointing it out. Please refer to the general response where we discuss how $\\lambda_1$ and $\\lambda_2$ should be reasonably small and do not dominate the error term in practice. We will add the discussions in the revision."}, {"review_id": "Nt4k2jNJp5", "reviewer": "Reviewer_6EPu", "summary": "The 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.", "questions": "##### 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\".", "limitations": "The authors have adequately and precisely addressed the limitations in this work.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "+ 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.", "weaknesses": "+ Please refer to the subsequent question section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback! We provide discussions on your concerns as follows.\n\nQ: **The paper would benefit from additional theoretical analysis on the intuition of similar optimal state distributions.**\n\nA: Thm. 4.3 (lines 258-259) of our paper provides evidence for the existence of such similar state distributions. We restate the theorem here:\n\nConsider two homomorphous MDPs with dynamics $T$ and $T^{\\prime}$. If $T^{\\prime} \\in\\left(T, \\varepsilon_m\\right)$ and the action gap $\\Delta$ follows $\\Delta>\\frac{(2-\\gamma) \\lambda_1 \\lambda_2 \\varepsilon_m}{1-\\gamma}$, for all $s \\in \\mathcal{S}$ we have $d_T^*(s)=d_{T^{\\prime}}^*(s)$.\n\nIn this theorem, $T$ and $T^{\\prime}$ refer to two distinct dynamics with similar structures. $d_T^*(s)$ and $d_{T^{\\prime}}^*(s)$ are optimal state distributions in these two environment dynamics. With the assumption of large enough action gap $\\Delta$ (which is mild as analyzed in line 260-265), it can be proved that $d_T^*(s)$ will be the same with $d_{T^{\\prime}}^*(s)$. This is in accordance with our initial intuition of similar optimal state distributions and the design of the SRPO algorithm. These discussions will also be added in the revision of our paper.\n\nQ: **It would be beneficial and interesting to provide more extensive analyses or explanations about the two types of regularization (DARA style and SRPO style) and demonstrate in what situations SRPO can perform larger advantage and why SRPO can work with more cases.**\n\nA:  Thanks for pointing it out. We will add the following analysis to the experiment part in the revision. According to Tab. 2 of our paper, SRPO emerges as the top-performing algorithm across all environments. However, in the Walker2d and HalfCheetah environments SRPO only exhibits a relatively small margin of improvement over the DARA style of algorithm (behavior regularizing). This observation can be attributed to the dynamics functions of Walker2d and HalfCheetah tasks being less affected by the altered environment parameters. In other words, these two tasks are easier to solve. Such task property is in turn substantiated by the performance of baseline algorithms without context encoders, such as GAIL and MOPO, which manage to achieve non-zero performance scores in these environments.\n\nIn contrast, in environments characterized by more significant dynamics deviations, such as Hopper, SRPO demonstrates a more pronounced margin of improvement over baseline algorithms. This heightened performance can be attributed to SRPO's focus on a more general property — the intuition of similar optimal state distributions — enabling it to effectively address the challenges posed by environments with larger dynamics deviations.\n\n\nQ: **Further elaboration is needed to support the statement that SRPO enables efficient data reuse.**\n\nA: Thanks for pointing it out. We will add the following elaborations to the experiment analysis in the revision. As for *online experiments*, Fig. 4 of our paper demonstrates that in 7 out of 8 experiment settings, SRPO consistently outperforms baseline algorithms throughout the training process, even with very few training steps. This observation highlights SRPO's ability to achieve good performance with limited training data in online settings, making it a more data-efficient approach.\n\nIn *offline experiments*, we reduce the amount of offline data to 1/3 and perform additional experiments. The results are shown in the following table, where M and ME refer to Medium and Medium-Expert performance levels of offline data. MAPLE+SRPO can still achieve better performance than baseline algorithms. **MAPLE+SRPO improves the performance by 31\\% over MAPLE+DARA and 49\\% over MAPLE**. These evidences indicate that SRPO indeed enables efficient data reuse, which is in accordance with statements in the introduction part.\n\n|                |      MOPO       |      MAPLE      |    MAPLE+DARA     | MAPLE+SRPO (ours)  |\n| -------------- | :-------------: | :-------------: | :---------------: | :----------------: |\n| Walker2d-M     | 0.21 $\\pm$ 0.13 | 0.45 $\\pm$ 0.18 |  0.74$\\pm$ 0.12   | **0.79**$\\pm$0.04  |\n| Walker2d-ME    |  0.14$\\pm$0.06  | 0.26 $\\pm$ 0.01 |  0.38$\\pm$ 0.03   | **0.61**$\\pm$ 0.11 |\n| Hopper-M       |  0.01$\\pm$0.00  |  0.42$\\pm$0.36  |   0.36$\\pm$0.06   | **0.51**$\\pm$0.14  |\n| Hopper-ME      |  0.01$\\pm$0.00  |  0.33$\\pm$0.09  |   0.16$\\pm$0.04   | **0.40**$\\pm$0.06  |\n| HalfCheetah-M  |  0.10$\\pm$0.01  |  0.50$\\pm$0.06  |   0.37$\\pm$0.01   | **0.55**$\\pm$0.03  |\n| HalfCheetah-ME | -0.03$\\pm$0.00  |  0.35$\\pm$0.01  | **0.63**$\\pm$0.03 |   0.62$\\pm$0.19    |\n| Average        |      0.07       |      0.39       |       0.44        |        **0.58**        |\n\nQ: **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\nA: We would like to emphasize that the training of the discriminator in our paper is inherently different from that in generative adversarial networks (GANs). Unlike GANs, which are notoriously challenging to stabilize during training, our approach does not involve a generator module. Consequently, the discriminator in our study serves as a simple classifier, leading to quick convergence and enhanced training stability, as shown in Fig. 2 of the one-page pdf. Therefore, it is unnecessary to take specific design considerations to stabilize the training of the discriminator. "}]

Author rebuttal:
We extend our sincere gratitude to all the reviewers for providing insightful comments. Both reviewer PXH4 and ZrWb mentioned the absence of discussions on the Lipstchitz coefficients in the assumptions. In this common response, we will address this issue and discuss how Lipstchitz coefficients can be small and negligible in practical scenarios. These discussions will be incorporated into the revised version of our paper.

We first restate our assumptions on the Lipschitz properties: The reward function $r(s, a, s^{\prime})$ w.r.t. the action $a$ is $\lambda_1$-Lipschitz and the dynamics function $T(s, a)$ w.r.t. the action $a$ is $\lambda_2$-inverse Lipschitz. According to Eq. 7 (line 241-242), two terms related to Lipschitz properties are considered: the multiplication $\lambda_1 \lambda_2$ and $\lambda_1$ itself.

Regarding the reward functions, the Lipschitz property implies that if $s$ and $s^{\prime}$ keeps unchanged, the deviation of the reward $r$ will be no larger than $\lambda_1$ times the deviation of the action $a$. Therefore, the Lipstchiz coefficient $\lambda_1$ is solely related to action-related terms in the reward function. It is important to note that different actions exist given the same $s$ and $s^{\prime}$ since we may compute the reward function in different dynamics. Considering the dynamics functions, the Lipschitz property indicates that if the current state $s$ remains unchanged and the actions differ by : $\|a_1-a_2\| \geqslant \varepsilon$, the next states will exhibit a significant difference: $\|s_1^{\prime}-s_2^{\prime}\| \geqslant \frac{\varepsilon}{\lambda_2}$.

In the following table, we list the action-related terms of the reward functions for various RL evaluation environments, along with the corresponding values of $\lambda_1$ derived from these terms. Additionally, we sample 50,000 (s,a,s') tuples from the replay buffer, slightly modify the action, and observe how the resulting next state $s^{\prime}$ changes. The replay buffer contains trajectories collected during different training phases, and should be diverse enough to cover most possible trajectories. This empirical analysis allows us to calculate $\lambda_2$ in practice and demonstrate it in the table.

| Environment | Action-related Reward | $\lambda_1$ | $\lambda_2$ | $R_{\max }$ |
| :--- | :--- | :--- | :--- | :--- |
| CartPole-v0 | 0 | 0 | 1.42 | 1.00 |
| InvertedPendulum-v2 | 0 | 0 | 8.58 | 1.00 |
| Swimmer-v2 | $-0.0001\|a\|_2^2$ | 0.0001 | 2.59 | 0.36 |
| HalfCheetah-v2 | $-0.1\|a\|_2^2$ | 0.1 | 1.01 | 4.80 |
| Hopper-v2 | $-0.001\|a\|_2^2$ | 0.001 | 3.45 | 3.80 |
| Walker2d-v2 | $-0.001\|a\|_2^2$ | 0.001 | 4.70 | $\geqslant 4$ |
| Ant-v2 | $-0.5\|a\|_2^2$ | 0.5 | 0.69 | 6.0 |
| Humanoid-v2 | $-0.1\|a\|_2^2$ | 0.1 | 0.03 | $\geqslant 8$ |

As indicated in the table, the action-related terms in reward functions exhibit reasonably small coefficients across all environments, leading to small $\lambda_1$ values. Combined with medium values of $\lambda_2$, it can be inferred that Lipstchitz terms, including $\lambda_1$ and $\lambda_1\lambda_2$,  will remain small in practical scenarios and will not dominate the error term in Eq. 7. Also, the action gap assumption in Thm. 4.3 (line 258-259) is not strong and holds in many situations.

Meta-review: The paper proposes a context-aware method called SRPO (State Regularized Policy Optimization) to address the issue of training RL algorithms on data with dynamics shifts. Paper is well-written and easy to follow, both theoretical analysis (lower-bound performance guarantees) and experiments are intuitive and strong. The setting studied is also general enough to be applicable to many real-world scenarios. Methods can also be applied to both online and offline RL settings.
Decision: Accept (poster)