Model-Based Offline Reinforcement Learning with Pessimism-Modulated Dynamics Belief

With the converged Q-value, we are ready to establish a more direct connection between the AMG and the original MDP. The connection appears as the answer to a natural question: the calculation of (6) encompasses biased samples from the dynamics belief distribution, can we treat these samples as the unbiased ones sampling from another belief distribution? We give positive answer in the following theorem.

In right-hand side of (7), $\tau^{sa}$ itself is random following the belief distribution, thus $\mathbb{E}_{\tau^{sa},\pi}\big{[}Q^{\pi}_{N,k}(s^{\prime},a^{\prime})\big{]}$, as a functional of $\tau^{sa}$, is also a random variable, whose cumulative density function is determined by the belief distribution $\mathbb{P}_{T}^{sa}$. Intuitively, we can treat $\mathbb{E}_{\tau^{sa},\pi}\big{[}Q^{\pi}_{N,k}(s^{\prime},a^{\prime})\big{]}$ as a pessimism indicator for transition $\tau^{sa}$, with larger value indicating less pessimism.

From Theorem 2, the maximum of $w$ is obtained at $\tau^{*}\!\!:\!F\!\left(\mathbb{E}_{\tau^{*},\pi}\big{[}Q^{\pi}_{N,k}(s^{\prime},a^{\prime})\big{]}\right)\!\!=\!\frac{k-1}{N-1}$, i.e., the transition with $\frac{k-1}{N-1}$-quantile pessimism indicator. Besides, when $\mathbb{E}_{\tau^{sa},\pi}\big{[}Q^{\pi}_{N,k}(s^{\prime},a^{\prime})\big{]}$ departs the $\frac{k-1}{N-1}$ quantile, the reweighting coefficient for its $\tau^{sa}$ decreases. Considering the effect of $w$ to $\widetilde{\mathbb{P}}_{T}^{sa}$ and the equivalence between the AMG and the refined MDP, we can say that $J(\pi)$ is a soft percentile performance. Compared to the standard percentile criteria, $J(\pi)$ is derived by reshaping belief distribution towards concentrating around a certain percentile, rather than fixating on a single percentile point. Due to this feature, we term $\widetilde{\mathbb{P}}_{T}^{sa}$ Pessimism-Modulated Dynamics Belief (PMDB).

Lastly, recall that all the above derivations are with hyperparameters $k$ and $N$, we present the monotonicity of $Q^{\pi}_{N,k}$ over them in Theorem 3. Furthermore, by combining Theorem 1 with Theorem 3, we conclude that $J(\pi)$ decreases with $N$ and increases with $k$.

Remark 1 (Special Cases). For $N=k=1$, we have $\widetilde{\mathbb{P}}^{sa}_{T}=\mathbb{P}_{T}^{sa}$. Then, the performance is evaluated through sampling the initial belief distribution. This resembles the common methodology in model-based RL (MBRL), with dynamics belief defined by the uniform distribution over dynamics model ensembles. For $k=\delta(N-1)+1$ and $N\rightarrow\infty$, $\widetilde{\mathbb{P}}_{T}^{sa}$ asymptotically collapses to be a delta function. Then, $J(\pi)$ degrades to fixate on a single transition instance. It is equivalent to the robust MDP with the uncertainty set constructed as $\Big{\{}\tau^{sa}:\mathbb{P}_{T}^{sa}(\tau^{sa})>0,\mathbb{E}_{\tau^{sa},\pi}\big{[}Q^{\pi}_{N,k}(s^{\prime},a^{\prime})\big{]}\geq F^{-1}(\delta)\Big{\}}$. In this sense, the AMG is a successive interpolation between MBRL and robust MDP.

Define the KL-regularized return for the AMG by

where $\alpha\geq 0$ is the strength of regularization, and $\mu$ is a reference policy to keep close with.

KL-regularized MDP is considered in previous works to enhance exploration, improve robustness to noise or insert expert knowledge SQL ; SAC ; RegularizedMDP ; ISKL ; Prior . Here, the idea is to constrain the optimized policy in neighbour of a reference policy so that the inner problem is adequately evaluated for such a small policy region.

To optimize $\bar{J}(\pi;\mu)$, we introduce the soft $(N,k)$-pessimistic Bellman backup operator:

Apparently, the solved policy $\bar{\pi}^{*}$ depends on the reference policy $\mu$, and setting $\mu$ arbitrarily aforehand can result in suboptimal policy for $J(\pi)$. In fact, we can construct a sequence of sub-problems, with $\mu$ chosen as an improved policy from last sub-problem. By successively solving them, the impact from the initial reference policy is gradually eliminated.

Ideally, by combining Theorems 4 and 5, the policy for $J(\pi)$ can be continuously improved by infinitely applying soft pessimistic Bellman backup operator for each of the sequential sub-problems.

Remark 2 (Iterative Regularized Policy Optimization as Expectation–Maximization with Structured Variational Posterior). According to Theorem 2, PMDB can be recovered with the converged Q-function $Q^{\pi^{*}}_{N,k}$. From an end-to-end view, we have an initial dynamics belief $\mathbb{P}_{T}^{sa}$, then via the calculation based on the belief samples and the reward function, we obtain the updated dynamics belief $\widetilde{\mathbb{P}}_{T}^{sa}$. It is likely that we are doing some form of posterior inference, where the evidence comes from the reward function. In fact, the iterative regularized policy optimization can be formally recast as an Expectation-Maximization algorithm for offline policy optimization, where the Expectation step correponds to a structured variational inference procedure for dynamics. We elaborate it in Appendix C.

While solving each sub-problem makes prominent update to policy compared with the policy gradient method, we may need to construct several sub-problems before convergence, then exactly solving each of them incurs unnecessary computation. For practical consideration, next we introduce a smooth-evolving reference policy, with which the explicit boundary between sub-problems is blurred. Based on this reference policy, and by further adopting function approximator, we devise an offline RL algorithm to approximately maximize $J(\pi)$.

The idea of smooth-evolving reference policy is inspired by the soft-updated target network in deep RL literature DQN ; DDPG . That is setting the reference policy as a slowly tracked copy of the policy being optimized. Formally, consider a parameterized policy $\pi_{\phi}$ with the parameter $\phi$. The reference policy is set as $\mu=\pi_{\phi^{\prime}}$, where $\phi^{\prime}$ is the moving average of $\phi:\phi^{\prime}\leftarrow\omega_{1}\phi+(1-\omega_{1})\phi^{\prime}$. With small enough $\omega_{1}$, the Q-value of the state-action pairs induced by $\pi_{\phi^{\prime}}$ (or its slight variant) can be sufficiently evaluated, before being used to update the policy. Next, we detail the loss functions to learn Q-value and policy with neural network approximators.

Denote the parameterized Q-function by $Q_{\theta}$ with the parameter $\theta$. It is trained by minimizing the Bellman residual of both the AMG and the empirical MDP:

with

where $Q_{\theta^{\prime}}$ represent the target Q-value softly updated for stability DDPG , i.e., $\theta^{\prime}\leftarrow\omega_{2}\theta+(1-\omega_{2})\theta^{\prime}$, and $\mathcal{D}^{\prime}$ is the on-policy data buffer for the AMG. Since the game transition is known, the game can be executed with multiple counterparts in parallel, and the buffer only collects the latest sample for each of them. To promote direct learning from $\mathcal{D}$, we also include the Bellman residual of the empirical MDP in (11).

As with policy update, Theorem 4 states that the optimal policy for $\bar{J}(\pi;\mu)$ is propotional to $\mu(a|s)\exp{\left(\frac{1}{\alpha}\bar{Q}^{*}_{N,k}(s,a)\right)}$. Then, we update $\pi_{\phi}$ by supervisedly learning this policy, with $\mu$ and $\bar{Q}^{*}_{N,k}$ replaced by the smooth-evolving reference policy and the learned Q-value:

where $A$ and $B$ are constant terms. In general, (4.2) can be replaced by any tractable function that measures the similarity of distributions. For example, when $\pi_{\phi}$ is Gaussian, we can apply the recent proposed $\beta$-NLL Pitfall , in which each data point’s contribution to the negative log-likelihood loss is weighted by the $\beta$-exponentiated variance to improve learning heteroscedastic behavior.

To summarize, the algorithm alternates between collecting on-policy data samples in AMG and updating the function approximators. In detail, the latter procedure includes updating the Q-value with (11), updating the optimized policy with (4.2), and updating the target Q-value as well as the reference policy with the moving-average rule. The complete algorithm is listed in Appendix D.

We compare the proposed offline RL algorithm with the baselines including: BEAR BEAR and BRAC BRAC , the model-free approaches based on policy constraint, CQL CQL , the model-free approach by penalizing Q-value, EDAC EDAC , the previous SoTA on the D4RL benchmark, MOReL MOREL , the model-based approach which terminates the trajectory if the dynamics uncertainty exceeds a certain degree, and BC, the behavior cloning method. These approaches are evaluated on a total of eighteen domains involving three environments (hopper, walker2d, halfcheetah) and six dataset types (random, medium, expert, medium-expert, medium-replay, full-replay) per environment. We use the v2 version of each dataset.

The results are summarized in Table LABEL:Table:D4RL. Our approach PMDB obviously improves over the previous SoTA on 9 tasks and performs competitively in the rest. Although EDAC achieves better performance in walker2d with several dataset types, its hyperparameters are tuned individually for each task. The later experiments on the impact of hyperparameters indicate that larger $k$ or smaller $N$ could generate better results for walker2d and halfcheetah. We also find that PMDB significantly outperforms MOReL, another model-based approach. It is encouraging that our model-based approach achieves competitive or better performance compared with the SoTA model-free approach, as model-based approach naturally has better support for multi-task learning and transfer learning, where the offline data from relevant tasks can be further leveraged.

Figure 2 presents the learning curves in the AMG, as well as the received return when deploying the policy being learned in true MDP. The performance in the AMG closely tracks the true performance from the lower side, implying that it can act as a reasonable surrogate to evaluate/optimize performance for the true MDP. Besides, the performance in the AMG improves nearly monotonously, verifying the effectiveness of the proposed algorithm to approximately solve the game.

Recall that PMDB does not explicitly quantify dynamics uncertainty to penalize return, Figure 3 checks how the dynamics uncertainty and the Q-value of visited state-action pairs change during the learning process. The uncertainty is measured by the logarithm of standard deviation of the predicted means from the $N$ dynamics samples, i.e., $\log\left(\text{std}\left(\mathbb{E}_{\tau}[s^{\prime}];\tau\in\mathcal{T}\right)\right)$. The policy being learned is periodically tested in the AMG for ten trials, and we collect the whole ten trajectories of state-action pairs. The solid curves in Figure 3 denote the mean uncertainty and Q-value over the collected pairs, and shaded regions denote the standard deviation. From the results, the dynamics uncertainty first sharply decreases and then keeps a slowly increasing trend. Besides, in the long-term view, the Q-value is correlated with the degree of uncertainty negatively in the first phase and positively in the second phase. This indicates that the policy first moves to the in-distribution region and then tries to get away by resorting to the generalization of dynamics model.

In Figure 3, we also notice that the large dip of Q-value is accompanied with the sudden raise of dynamics uncertainty. We suspect this is due to the optimized policy being far from the dataset. We try to verify by checking the maximal return covered by the offline dataset. It shows that the maximal normalized returns provided by offline datasets are 99.6, 92.0 and 45.0 respectively for hopper, walker2d and halfcheetah, while the proposed approach achieves 106.8, 94.2 and 75.6. The policy optimization is more significant for halfcheetah (where we observe the large dip), indicating the policy should move further from the dataset.

The above finding also explains why the AMG performance in Figure 2 runs into large dip only for halfcheetah: with the larger dynamics uncertainty, the secondary player can choose the more pessimistic transition. However, we want to highlight that it is normal behavior of the proposed algorithm and does not mean instability, as we are handling the alternating Markov game, a specialization of zero-sum game. Besides, we can see that even when the AMG performance goes down the MDP performance is still stable.

Table 2 lists the impact of $k$. In each setting, we evaluate the learned policy in both the true MDP and the AMG. The performance in the AMGs improve when increasing $k$. This is consistent with the theoretical result, even that we approximately solve the game. Regarding the performance in true MDPs, we notice that $k=2$ corresponds to the best performance for hopper, but for the others $k=3$ is better. This indicates that tuning hyperparameter online can further improve the performance. The impact of $N$ is presented in Appendix H, suggesting the opposite monotonicity.