SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhenghai Xue $^{1}$  Qingpeng Cai $^{2}$  Shuchang Liu $^{2}$  Dong Zheng $^{2}$  Peng Jiang $^{2}$  Kun Gai $^{3}$  Bo An $^{1}$
$^{1}$ Nanyang Technology University, Singapore
$^{2}$ Kuaishou Technology  $^{3}$  Unaffiliated
zhenghai001@e.ntu.edu.sg boan@ntu.edu.sg gai.kun@qq.com
{caiqingpeng, liushuchang, zhengdong, jiangpeng}@kuaishou.com
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 (State Regularized Policy Optimization) 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.
Reinforcement Learning (RL) has achieved great success in solving challenging sequential decision-making 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
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 (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] 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.
A Markov Decision Process (MDP) can be defined by a tuple  $(S, \mathcal{A}, T, r, \gamma, \rho_0)$ , where  $\mathcal{S}$  is the state space,  $\mathcal{A}$  is the bounded action space with actions  $a \in (-1, 1)$ ,  $T(s'|s, a) \in [0, 1]$  and  $r(s, a, s') \in [-R_{\max}, R_{\max}]$  are the transition and reward functions.  $\gamma \in (0, 1)$  is the discount factor and  $\rho_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, \varepsilon)$  as  $\{T' \mid |T(s, a) - T'(s, a)| < \varepsilon, \forall s \in S, a \in \mathcal{A}\}$  which is the  $\varepsilon$ -neighbourhood of  $T$ . RL aims at maximizing the accumulated return of policy  $\pi$ :
Figure 1: Performance comparison of PPO [7], CaDM [9] and  $\mathrm{CaDM + SRPO}$  in the Ant environment, where SRPO is our proposed state regularized policy optimization method. Details of the experiment setup are in Sec. 5.1.
$\eta_T(\pi) = E_{\pi, T}\left[\sum_{t=0}^{\infty} \gamma^t r(s_t, a_t)\right]$ , where the expectation is computed with  $s_0 \sim \rho_0$ ,  $a_t \sim \pi(\cdot|s_t)$ , and  $s_{t+1} \sim T(\cdot|s_t, a_t)$ . The optimal policy  $\pi^*$  is defined as  $\pi_T^* = \arg \max_{\pi} \eta_T(\pi)$ . In an MDP with a policy  $\pi$ , the Q-value  $Q_T^\pi(s, a)$  denotes the expected return after taking action  $a$  at state  $s$ :  $Q_T^\pi(s, a) = E_{\pi, T}\left[\sum_{t=0}^{\infty} \gamma^t r(s_t, a_t) | s_0 = s, a_0 = a\right]$ . The value function is defined as  $V_T^\pi(s) = \mathbb{E}_{a \sim \pi(\cdot|s)} Q_T^\pi(s, a)$  with  $V_T^*(s)$  being the shorthand for  $V_T^{\pi_T^*}(s)$ . It satisfies the optimal Bellman Equation  $V_T^*(s) = \max_{a} r(s, a) + \gamma \mathbb{E}_{s' \sim T(\cdot|s, a)} V_T^*(s')$ . We can also define the stationary state distribution (also known as state occupation function) as  $d_T^\pi(s) := (1 - \gamma) \sum_{t=0}^{\infty} \gamma^t P_T(s_t = s \mid \pi)$  with  $d_T^*(s)$  being the shorthand for  $d_T^{\pi_T^*}(s)$ .
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, \mathcal{A}, \Theta, T, r, \gamma, \rho_0)$ , where  $\Theta$  is the space of hidden parameters. The transition function  $T_{\theta}(s'|s, a, \theta)$  is parameterized not only by states and actions, but also by a hidden parameter  $\theta$  sampled from  $\Theta$ . The action gap of an HiP-MDP is defined as  $\Delta = \min_{\theta \in \Theta} \min_{s \in S} \min_{a \neq \pi^*(s)} V_{T_\theta}^*(s) - Q_{T_\theta}^*(s, a)$ , which reflects the minimum gap between an optimal action and all other sub-optimal actions.
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 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 [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:
$$
\max  _ {\pi} \mathbb {E} _ {s, a \sim \mathcal {D}} \left[ \mathbb {E} _ {a ^ {\prime} \sim \pi (\cdot | s)} Q (s, a ^ {\prime}) \right] \quad \text {s . t .} \quad \mathbb {E} _ {s \sim \mathcal {D}} \left[ \hat {D} (\pi (\cdot | s) \| \hat {\pi} (\cdot | s)) \right] <   \varepsilon , \tag {1}
$$
where  $\mathcal{D}$  is the replay buffer,  $\hat{\pi}$  is the regularizing policy, and  $\hat{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,  $\hat{\pi}$  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^\pi (s)$  of policy  $\pi$  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.
Figure 2: Visualization of state and action densities in data sampled from the Inverted Pendulum environment with gravity 5 and 10. Under both gravities, 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.
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:
$$
\max  _ {\pi} \mathbb {E} _ {s _ {t}, a _ {t} \sim \tau_ {\pi}} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r \left(s _ {t}, a _ {t}\right) \right] \quad \text {s . t .} D _ {\mathrm {K L}} \left(d _ {\pi} (\cdot) \| \zeta (\cdot)\right) <   \varepsilon , \tag {2}
$$
where  $\zeta(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<sup>1</sup>.
We solve Eq. (2) by casting it to the following unconstrained optimization problem via Lagrange multipliers<sup>2</sup>:
$$
L = - \mathbb {E} _ {s _ {t}, a _ {t} \sim \tau} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \left(r (s _ {t}, a _ {t}) + \lambda \log \frac {\zeta (s _ {t})}{d _ {\pi} (s _ {t})}\right) \right] - \frac {\lambda \varepsilon}{1 - \gamma}, \tag {3}
$$
where  $\lambda > 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  $\lambda \log \frac{\zeta(s_t)}{d_\pi(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  $\frac{\zeta(s)}{d_{\pi}(s)}$  because obtaining the optimal state distribution during online training or given suboptimal offline dataset is infeasible. Also,  $d_{\pi}(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  $\frac{\zeta(s)}{d_{\pi}(s)}$ .
Proposition 3.1. In a GAN, when the real data distribution is  $\zeta(s)$  and the generated data distribution is  $d_{\pi}(s)$ , the output of the discriminator  $D(s)$  follows
$$
\frac {D (s)}{1 - D (s)} = \frac {\zeta (s)}{d _ {\pi} (s)}. \tag {4}
$$
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  $\zeta(s)$  and away from  $d_{\pi}(s)$ , which is sort of the average state distribution. Motivated by [16], we model state optimality by a variable  $\mathcal{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  $\mathcal{O}_t$ .  $\mathcal{O}_t = 1$  denotes that  $s_t$  is the optimal state at timestep  $t$ . The observation model is given by
Figure 3: HMM in MDP with optimality variables  $\mathcal{O}_t$ .
$$
p \left(\mathcal {O} _ {t} \mid s _ {t}\right) = \max  _ {a _ {t}} \exp \left[ \gamma^ {t} \left(r \left(s _ {t}, a _ {t}\right) - R _ {\max }\right) \right]. \tag {5}
$$
We can therefore compute the state density ratio  $\frac{\zeta(s)}{d_{\pi}(s)}$  as
$$
\frac {\zeta (s)}{d _ {\pi} (s)} = \frac {d _ {\pi} (s | \mathcal {O} _ {0 : \infty})}{d _ {\pi} (s)} = \frac {p (\mathcal {O} _ {0 : \infty} | s , \pi) d _ {\pi} (s)}{p (\mathcal {O} _ {0 : \infty} | \pi) d _ {\pi} (s)} = \frac {\mathbb {E} _ {t} [ p (\mathcal {O} _ {0 : t - 1} | s _ {t} , \pi) p (\mathcal {O} _ {t : \infty} | s _ {t} , \pi) ]}{p (\mathcal {O} _ {0 : \infty} | \pi)}, \tag {6}
$$
where the second equation follows the Bayes' law. The last term is related to the forward probability  $\alpha_{t}(s_{t}) = p(\mathcal{O}_{0:t - 1}|s_{t},\pi)$  and backward probability  $\beta_{t}(s_{t}) = p(\mathcal{O}_{t:\infty}|s_{t},\pi)$  in the HMM. We discuss in Appendix A.2 that  $\beta_{t}(s_{t})$  is positively related to a soft version of MDP's state value  $V_{\pi}(s_t)$ . Also,  $\alpha_{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  $\zeta (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)$ .
<sup>1</sup>A similar form of Eq. (2) (See Eq. 1 in Sec. 2.2) is employed in Offline RL algorithms to ensure conservative policy updates. But it restricts the training data to be sampled from the same environment.
2The detailed derivations of the Lagrangian can be found in Appendix A.1.
Algorithm 1 The workflow of SRPO on top of MAPLE [12].
1: Input:  $\phi_{\varphi}$  as a context encoder parameterized by  $\varphi$ ; Adaptable policy network  $\pi_{\theta}$  parameterized by  $\theta$ ; Adaptable value network  $V_{\psi}$  parameterized by  $\psi$ ; Offline dataset  $\mathcal{D}_{\mathrm{off}}$ ; Rollout horizon  $H$ ; State partition ratio  $\rho$ ; State discriminator  $D_{\delta}$  parameterized by  $\delta$ ; Regularization coefficient  $\lambda$ .
2: for 1, 2, 3, ... do
3: for  $t = 1,2,\dots ,H$  do
4: Sample  $z_{t}$  from  $\phi_{\varphi}(z \mid s_{t}, a_{t-1}, z_{t-1})$  and then sample  $a_{t}$  from  $\pi_{\theta}(a \mid 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  $\mathcal{D}_{\mathrm{rollout}}$
6: end for
7: Update the context encoder  $\phi_{\varphi}$  according to MAPLE.
8: Sample a batch  $\mathcal{D}_{\mathrm{batch}}$  from  $\mathcal{D}_{\mathrm{off}}$  and  $\mathcal{D}_{\mathrm{rollout}}$  and rank them by their state-values estimated by  $V_{\psi}$ ; Add  $\rho |\mathcal{D}_{\mathrm{batch}}|$  states with higher state-values to  $\mathcal{D}_{\mathrm{real}}$  and the others to  $\mathcal{D}_{\mathrm{fake}}$ .
9: Train the discriminator  $D_{\delta}$  with nll loss.
10: For one-step transition  $(s_{t + 1},r_t,d_{t + 1},s_t,a_t,z_t)$  in  $\mathcal{D}_{\mathrm{batch}}$  , update  $r_t$  with  $r_t + \lambda \frac{D_\delta(s_t)}{1 - D_\delta(s_t)}$
11: Use the updated  $\mathcal{D}_{\mathrm{batch}}$  and SAC to update the policy and value network parameters  $\theta$  and  $\psi$ .
12: end for
3.4 Practical Algorithm
Summarizing the previous derivations, we obtain a practical reward regularization algorithm, termed as SRPO (State Regularized Policy Optimization) 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  $\rho$  states with higher rewards and add them to the dataset  $D_{\mathrm{real}}$ .  $D_{\mathrm{fake}}$  is similarly generated by states with lower rewards (line 10). We set  $\rho = 0.5$  in offline experiments with medium-expert level of data.  $\rho = 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  $\lambda \log \frac{\zeta(s)}{d_{\pi}(s)}$ , which is added to the reward  $r_t$  (line 12).  $\lambda$  is regarded as a hyperparameter with values 0.1 or 0.3. The effect of  $\lambda$  is investigated in Sec. 5.3. The procedure of the online algorithm CaDM [9]+SRPO is similar to MAPLE+SRPO, where the datasets  $D_{\mathrm{real}}$  and  $D_{\mathrm{fake}}$  are built with data from the replay buffer, rather than the offline dataset.