DreamSmooth: Improving Model-based Reinforcement Learning via Reward Smoothing

Author response

We appreciate the reviewer’s thorough and constructive feedback about our paper. We provide detailed answers to the questions below.   **[Q] Do the authors conduct some experiments with different prediction function head to justify whether it could solve the sparse rewards problem?** This is a really good point! As suggested by the reviewer, we conducted additional experiments to see whether different reward head type (regression) and loss (L1 and L2) can solve the reward prediction problem. In our experiments on RoboDesk and Hand, both L1 and L2 losses with a typical MLP reward head miss many sparse reward signals, similar to the results with DreamerV3’s two-hot categorical regression. Furthermore, these reward modeling approaches also showed improved performances with DreamSmooth. We summarized these results in the updated paper (Figure 17, Appendix E). In summary, we found that the reward prediction problem exists in diverse commonly used regression models and losses, and DreamSmooth effectively alleviates this problem.   **[Q] why DreamSmooth works comparable without reward smooth technique in dense reward environments.** We would like to refer the reviewer to Theorem A.1 in Appendix A, which shows that the proposed EMA reward smoothing guarantees an optimal policy even with the changes in reward distribution. Thus, smoothing the past rewards does not deteriorate the performance of MBRL. In the case when reward smoothing involves future rewards, Theorem A.2 explains that an optimal policy under any of our reward smoothing functions is also optimal under the original reward function. This still does not guarantee that MBRL can achieve optimal performance with reward smoothing; but, in practice, blurring reward signals a few steps before and after does not completely change learning dynamics, especially under the dense reward setting.   ---   We thank the reviewer again for the time and effort put into improving our paper. Please feel free to let us know if there are any additional concerns or questions.

Author response

We appreciate the reviewer’s thorough and constructive feedback about our paper. We provide detailed answers to the questions below.   **[Q] Have you experimented with different loss functions (on the unsmoothed rewards)?** This is a really good point! As suggested by the reviewer, we conducted additional experiments to see whether different reward head type (regression) and loss (L1 and L2) can solve the reward prediction problem. In our experiments on RoboDesk, both L1 and L2 losses with a typical MLP reward head miss many sparse reward signals, similar to the results with DreamerV3’s two-hot categorical regression. Furthermore, these reward modeling approaches also showed improved performances with DreamSmooth. We summarized these results in the updated paper (Figure 17, Appendix E). In summary, we found that the reward prediction problem exists in diverse commonly used regression models and losses, and DreamSmooth effectively alleviates this problem.   **[Q] The need to experiment with 3 different reward smoothing schemes each with its own hyperparameters** Thank you for pointing this out. We agree that it could be more convincing with the use of consistent hyperparameters across tasks. We have updated the results (Figure 8) in our experiments with the same hyperparameters across all tasks and the hyperparameters can be found in Table 1: $\sigma=3$ for Gaussian smoothing, $\alpha=0.3, 0.33$ for EMA smoothing, and $\delta=9$ for Uniform smoothing. Our reward smoothing methods show consistent results across tasks and we could observe that the effect of the smoothing hyperparameters are marginal.   **[Q] The rationale behind why the counter-intuitive results in crafter is not convincing.** We have extended our analysis on Crafter in Appendix G. After carefully inspecting reward prediction of unsmoothed and smoothed rewards on Crafter, we found that Gaussian and uniform smoothing results in **more false positives** while achieving better prediction accuracy, as shown in Figure 20. We suspect this could be due to predicting rewards for the unknown future states. Asymmetric variants of the uniform smoothing kernel (Uniform [-4, 0] and Uniform [0, 4]) in Figure 21 clearly show that reward smoothing kernels dependent on future rewards are prone to predicting positive rewards even when the ground truth reward is 0. DreamSmooth’s high rate of false positives can be one reason for the poor performance on Crafter despite its high precision and recall on reward prediction. It is still surprising how DreamerV3 could perform well while missing most of the reward predictions. We leave this investigation as future work.   ---   We thank the reviewer again for the time and effort put into improving our paper. Please feel free to let us know if there are any additional concerns or questions.

Author response

We appreciate the reviewer’s thorough and constructive feedback about our paper. We provide detailed answers to the questions below.   **[Q] the fundamental challenge that comes with sparse rewards is exploration. … it does not address the fundamental exploration problem.** As the reviewer pointed out, our paper is **not** tackling all the challenges in solving sparse-reward tasks. Instead, we try to bring attention to one of the challenges in sparse-reward tasks, **the reward prediction problem** in model-based RL, which can happen even when exploration is not very difficult. For example, vanilla DreamerV3 could get sparse rewards in all our experiments but fails to correctly predict these sparse rewards and ends up with poor final performance. We made this point clearer in the abstract and introduction of the revised paper by removing some “sparse reward”.   **[Q] reward smoothing can give rise to overoptimistic behaviors. Do you see this reflected in any task?** Our reward smoothing will assign some positive reward to a bad state after a good state. This is not a big problem because of the following reasons. * After reward smoothing, the bad state (A) after the good state (B) can be assigned with some positive reward. But, if there is a better state (C) after (B) then (C) will get a larger reward than (A) after reward smoothing. Thus, an RL policy will learn to move toward the state with relatively larger original rewards, which will be also larger after smoothing. * As the reviewer mentioned, the recurrent state space model used in Dreamer helps discern bad states followed after good states. In the staircase example, the rewards will be bumped up right after falling off the staircase but this reward doesn’t last long since the past history tells an agent that it is far away from a rewarding state. Many of our experiments reflect this property: once an agent gets a sparse reward, coming back to the same state does not give additional rewards.   **[Q] How does reward ambiguity affect MBRL methods? … can you run a Markovian MBRL method like MBPO with reward smoothing and see if there's any improvement there?** Thank you for an interesting suggestion! We conducted additional experiments with a Markovian model-based RL method, MBPO [1], on the state-based environments, Hand and RoboDesk. Unfortunately, MBPO failed to learn any of these tasks, so we could not verify whether reward smoothing improves the performance or not. We will further investigate this in an easier environment in the camera-ready version. On the dense-reward DMC tasks, DreamSmooth does not hurt the Markovian MBRL method, MBPO, as shown in Appendix, Section F and Figure 18.   **[Q] Does reward smoothing benefit model-free methods?** The benefit of our reward smoothing approach is easier reward model learning. Thus, our approach does not help model-free methods, which do not learn to predict rewards.   ---   We hope our response addresses all your concerns and questions. Please let us know if there are any additional concerns or questions.   [1] Janner et al. “When to Trust Your Model: Model-Based Policy Optimization” NeurIPS 2019

Author response

We appreciate the reviewer’s thorough and constructive feedback about our paper. We provide detailed answers to the questions below.   **[Q] This is more like an engineering trick…** We would like to highlight that the major contribution of this paper is to bring attention to reward modeling in MBRL, which has been largely overlooked and not actively discussed in the field. We provide empirical evidence that shows the importance of accurate reward modeling and the effectiveness of our simple approach. In deep learning, many engineering tricks have been playing a critical role despite their simplicity. We believe it’s worth discussing such engineering tricks, which suggests a new perspective or problem to the community; and we believe our paper does the same.   **[Q] Is there any rationale for using DreamerV3 as the baseline for most tasks?** We mainly used DreamerV3 since it is the **only** model-based RL method that consistently works on a variety of tasks with high-dimensional inputs, sparse rewards, and long episode lengths. In Figure 9 of the original submission, TD-MPC could learn the Hand task and TD-MPC w/ DreamSmooth significantly improved TD-MPC. But, TD-MPC failed to learn other sparse-reward tasks. We also conducted additional experiments with another model-based RL method, MBPO [1], on the state-based environments, Hand and RoboDesk. However, MBPO couldn’t learn any of these tasks. We included these results in Section F and Figure 18.   **[Q] I didn’t see justifications for the correctness of this approach. … would a smooth reward function blur out the true critical states, so that the optimal policy does not visit the critical states?** We would like to refer the reviewer to Theorem A.1 in Appendix A, which shows that the proposed EMA reward smoothing guarantees an optimal policy. This requires a policy to access the history of states and the recurrent state structures used in many model-based RL methods enable a policy under smoothed rewards to hold Theorem A.1. In the case when reward smoothing involves future rewards, Theorem A.2 explains that an optimal policy under any of our reward smoothing functions is also optimal under the original reward function. This still does not guarantee that MBRL can achieve optimal performance with reward smoothing; but, in practice, blurring reward signals a few steps before and after does not completely change learning dynamics since a policy learns from a value function (or MC value estimate in TD-MPC), which already blurs out the rewards on the critical states.   **[Q] if this work is indeed only constrained to robotic tasks, I believe the authors need to make that clear in the paper…** This work is *not* specifically designed for the robotic domain. Our paper includes the most popular RL benchmarks, Atari and DMC. Moreover, Crafter and Atari benchmarks in our experiments are not robotic tasks. We used many robotic environments since these tasks have challenging properties of partial observability and sparse reward, which makes reward prediction difficult. We are happy to try our method on other non-robotic RL benchmarks with similar properties.   ---   We hope our response addresses all your concerns and questions. Please let us know if there are any additional concerns or questions.   [1] Janner et al. “When to Trust Your Model: Model-Based Policy Optimization” NeurIPS 2019