Maximum Entropy Model Correction in Reinforcement Learning

We thank the reviewer for their positive review of the paper.

I am unsure if the proposed method is scalable.

In terms of scalability and the cost of optimization procedure in MaxEnt MoCo, we point out three properties that we believe makes the procedure scalable:

1. Scale of the dual problem: The dual problem of optimization problems (P1) and (P2) has only $d$ parameters. This means that it does not scale with the size of MDP. The value of $d$ will be kept rather small, since the agent needs to store $d$ functions for $\phi_i$ and $\psi_i$. Moreover, the dual problem of (P2), which is the version implemented in practice, is $\beta^2/4$-strongly convex (See equation D.15). The convexity allows efficiently solving the problem.

2. Parallelization: Sampling from the model is usually done in batches. Correcting each of these samples is an independent optimization problem that can be parallelized with careful implementation.

3. Adjusting the computation cost: It is possible to adjust the amount of correction to meet the computational budget. Dual parameters equal to zero correspond to no correction and using the model itself. We expect that partially optimizing the dual parameters and making few updates on them still partially corrects the model. Another technique is to use stronger regularization of dual parameters by choosing a larger value for the hyper parameter $\beta$, which makes the loss more strongly convex and easier to optimize. We provide an analysis of MaxEnt MoCo for the general choice of $\beta$ in Theorem 5 and 6.

We have included the runtimes of our experiments in the new revision of the paper. We have also provided the computation time for the correction procedure, which is 0.2 - 1.5 seconds for all 144 state-actions depending on the setting.

Question 1:

The main reason for discarding $\phi_{1:k}$ and $\psi_{1:k}$ is solely to conserve memory. Otherwise, the required memory of the algorithm would grow by each iteration as $O(k)$. One other minor reason is that when the number of basis functions grows, the second term in the optimization problem (P2) will dominate the first term. This can be fixed by adjusting the parameter $\beta$ accordingly, but it would make the algorithm more complicated.

Question 2:

No, we never construct any value function as a linear combination of the basis functions. The performance of the algorithm does depend on how well the true value function is approximated with a linear combination of the basis functions. This is only for the purpose of theoretical analysis. The distance between the true value function and the span of basis functions is not necessarily zero and appears in our bounds. More importantly, we add the new value function obtained from planning to the basis functions which is different from the true value function. Yet, such issues can be avoided with a properly designed BasisCreation function in line 9 of the MoCoDyna algorithm, for example by subtracting the part of the new value function that belongs to the span of previous basis functions.

Question 3:

If the queries are accurate, the correction of MaxEnt MoCo and the convergence rate of MoCoVI improve with a larger value of $d$. However, the main benefits of our methods are achievable even with $d=1$. For example, our theory shows that if the model is accurate enough, MoCoVI can converge to the true value function and do so faster than VI. This holds even for $d=1$ and large/continuous MDPs. Also, a single constant basis function for $d=1$ might be effective in MaxEnt MoCo (See proposition 1 and lemma 1).

In large continuous MDPs used in deep RL, the query result $\psi_i$ will probably be obtained by training a neural network to estimate $E[\phi(X’_i)]$ at next state $X_i’$ from the input $(X_i, A_i)$ as a regression problem. This means that $d$ will determine the number of heads of a neural network (if a shared network is used for all $\phi_i$ functions) or the number of networks (if a separate network is trained for each $\psi_i$). The choice of $d$ is a design decision. It might be beneficial to use the computation resources to train a better/larger model instead of training a huge network for queries and doing a high amount of computation for the correction procedure. For these reasons, we expect the optimal choice of $d$ to be small even in large environments.

The paragraph below equation (3.2): I do not understand the following sentence: $\bar P$ is not constructed by the agent.

It means that the agent does not store $\bar P$ as a function of state-actions to next-state distributions the same way it does for $\hat P$. For example if the model $\hat P$ is a neural network, we do not create a new network for $\bar P$. Instead, we compute the distribution $\bar P(\cdot | x,a)$ when needed by solving the optimization problem (P1) or (P2).

We thank the reviewer for their positive review of the paper.

The main weakness in my view is the limited number of experiments, moreover the experiments are relatively simple. I would have liked to see how the approach is used and how its performs (also in terms of time) for large or continuous environments (which seems to be planned for future work).

We agree with the reviewer that more extensive experiments can be very insightful for our proposed approach.

I'm not quite sure about the related work and comparison to approaches that use multiple models to improve predictions, which would be a more appropriate comparison (apart from OS-VI, eg residual models or even ensemble models), along with a comparison of computational costs (during training or inference). Even the additional empirical results in the appendix are not particularly large.

Surely there’s a large body of research on dealing with model error in MBRL. Many of these efforts can be considered ways to improve the approximate dynamics $\hat P$. Consequently our method can be implemented along with them. The empirical comparison of the gains achieved with these techniques to our method is more reasonable in deep RL experiments and is to be studied in future work.

In the introduction, you say MaxEnt MoCo first obtains $E[\phi(X’)]$ for all $(x,a)$. That sounds expensive and seems to be computed on demand; but what is the cost of solving the lazy approximation (P1).

In large or continuous MDPs, we can use function approximation to estimate these expectations. It is equivalent of fitting a regression model $\psi$ to a dataset $(X_i, A_i, \phi(X_i’))$ where $(X_i, A_i)$ is the input and $\phi(X_i’)$ is the output. Then an estimate of that expectation can be obtained at any $(x,a)$ by simply evaluating $\psi$ on $(x,a)$.

Can you comment about computational cost and scalablity of the approach given that the experiments were limited?

In terms of scalability and the cost of optimization procedure in MaxEnt MoCo, we point out three properties that we believe makes the procedure scalable:

1. Scale of the dual problem: The dual problem of optimization problems (P1) and (P2) has only $d$ parameters. This means that it does not scale with the size of MDP. The value of $d$ will be kept rather small, since the agent needs to store $d$ functions for $\phi_i$ and $\psi_i$. Moreover, the dual problem of (P2), which is the version implemented in practice, is $\beta^2/4$-strongly convex (See equation D.15). The convexity allows efficiently solving the problem.

2. Parallelization: Sampling from the model is usually done in batches. Correcting each of these samples is an independent optimization problem that can be parallelized with careful implementation.

3. Adjusting the computation cost: It is possible to adjust the amount of correction to meet the computational budget. Dual parameters equal to zero correspond to no correction and using the model itself. We expect that partially optimizing the dual parameters and making few updates on them still partially corrects the model. Another technique is to use stronger regularization of dual parameters by choosing a larger value for the hyper parameter $\beta$, which makes the loss more strongly convex and easier to optimize. We provide an analysis of MaxEnt MoCo for the general choice of $\beta$ in Theorem 5 and 6.

We have included the runtimes of our experiments in the new revision of the paper. We have also provided the computation time for the correction procedure, which is 0.2 - 1.5 seconds for all 144 state-actions depending on the setting.

I may have missed it but would it make sense to compare how close the dynamics of the real environment with the predicted and corrected one? If that's the case, evaluations of how well the environment can be predicted might be helpful even without additional results from RL experiments.

We thank the reviewer for the suggestion. We have now included Figure 5 in the appendix to show the error of original and corrected dynamics in our experiments. It can be observed that the correction procedure successfully reduces the model error. The effect is stronger with larger values of $d$.

We want to thank the reviewer for their thorough review and feedback. Here are some clarifications about the points raised in the review.

1)

The reviewer has pointed out an insightful special case that shows why MoCoVI converges faster than VI. Assume that the model is perfect ($\epsilon_\text{Model} = 0$). Also for simplicity assume that both VI and MoCoVI obtain $PV$ for any value function $V$ perfectly ($\epsilon_\text{Query} = 0$). In this case, MoCoVI converges immediately and $V_0$ will be the solution ($V^{\pi_\text{PE}}$ for PE and $V^*$ for control). This is reflected in our bound in Theorem 2 as $\gamma’ = 0$.

To see this, observe that the correction procedure will not change $\hat P$ at all because distribution $\hat P(\cdot | x,a)$ satisfies the constraints in optimization problem (P1) and achieves the optimal objective value of $0$. Consequently, the corrected dynamics $\bar P$ used to calculate $V_0$ will be the model $\hat P$ itself, which due to our assumption of perfect model is equal to $P$. Then as described in the second paragraph of Section 4, we have $V_0 = V^*(R, \bar P) = V^*(R, \hat P) = V^*(R, P) = V^*$.

In general and regarding Theorem 2, we want to clarify that the big O notation used in the discussion of Theorem does not hide the constants $3 \sqrt{2}$ and $1/(1-\gamma)$. In fact, they do impact the asymptotic rate $O(\gamma’^k)$. This notation hides any constant coefficient for the terms in the bound (e.g. $1 - 3c_1 \lVert \epsilon_\text{Model} \rVert$ in the second inequality) but constants inside $\gamma’$ change the asymptotic rate similar to how $O(2^x)$ and $O(3^x)$ are different. The key property of this rate regardless of these constants is that as $\epsilon_\text{Model}$ goes to zero, the rate $\gamma’$ goes to zero as well. Thus, if the model error is smaller than a threshold, which is at most $(1-\gamma)/(3\sqrt{2})$, we will have $\gamma’ < \gamma$, and MoCoVI converges faster than VI. The aforementioned constants affect the value of this threshold, but not the existence of it.

The maximum term in $\gamma’$ captures how well the linear combination of the past calculated value functions in MoCoVI can approximate the true value function. This term is upper bounded by 1, so the theorem would still hold without it. However, it could be very small or potentially zero, if for example $V^*$ is well approximate with the initial basis functions. A tighter bound would be to use the geometric mean instead of the maximum, but we decided on this version for the sake of simplicity.

2)

We agree with the reviewer on this point. Though, the experimental setup for a model-based RL algorithm on a continuous MDP requires the use of function approximators such as DNN. Implementations that work competitively in practice usually involve a lot of implementation-level fine-tuning. In this paper we focus on the fundamental understanding of MaxEnt MoCo through theory and do not get involved with these implementation nuances. Yet, we believe the study of our algorithms in continuous MDPs with function approximation is a very interesting future work.

3)

This might be a misunderstanding as MaxEnt MoCo does not require sampling the dynamics $d$ times. Having $d$ constraints in (P1) or terms in the summation in (P2) may incorrectly suggest that we need the number of samples to be proportional to $d$ to obtain $\psi_i$ functions (which estimate $P\phi_i$). But this is not needed as the same set of samples can be used to find all $\psi_i$ for $i = 1, \dotsc, d$. For example, if we have a single sample from $X’ \sim \mathcal{P}(\cdot|x,a)$, we can form $\psi_i(x,a) = \phi_i(X’)$ for all $i = 1, \dotsc, d$.

As we mentioned in the beginning of Section 5 and 3.2, finding functions $\psi_i \colon \mathcal{X} \times \mathcal{A} \to \mathbb{R}$ is a regression problem. When using function approximation in continuous MDPs, this can be done by fitting a regression model to any dataset of real samples even though it doesn’t cover all state-action pairs. In finite MDPs and without the generalization of function approximation, any RL algorithm requires knowledge of all state-action pairs relevant to the problem. Still, the requirement of samples does not depend on the order $d$ of the algorithm as the same samples can be used for finding all query result functions $\psi_i$.