Paper: DreamSmooth: Improving Model-based Reinforcement Learning via Reward Smoothing (ICLR-2024)

Reviewer comment (Evaluation):
Counter-intuitive results in Crafter, where the method performs worse even after having a much better reward model.

Author response:
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.

Reviewer comment (Experiments):
The need to experiment with 3 different reward smoothing schemes each with its own hyperparameters (since none of them seems to consistent favourite across tasks).

Author response:
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.

Reviewer comment (Experiments):
Have you experimented with different loss functions (on the unsmoothed rewards)? For example, what would happen if you use an L1 Loss instead of an L2 Loss commonly used in literature?

Author response:
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.

Reviewer comment (Evaluation):
The rationale behind why the counter-intuitive results in crafter is not convincing. Did the authors perform further empirical studies / analysis?

Author response:
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.

Reviewer comment (Theory):
While reward prediction error is indeed one consequence of reward sparsity, the fundamental challenge that comes with sparse rewards is exploration. If there is no reward signal in the first place, then reward smoothing does not work either. While this paper provides a simple remedy to alleviate reward sparsity, it does not address the fundamental exploration problem.

Author response:
**[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”.

Reviewer comment (Experiments):
It seems that reward smoothing can potentially lead to reward ambiguity. For example, if a bad state is visited right after a successful state vs. after a sequence of bad states, it would get assigned different smoothed reward values. How does reward ambiguity affect MBRL methods? I suspect the recurrent architecture of Dreamer helps mitigate this issue. To verify, can you run a Markovian MBRL method like MBPO with reward smoothing and see if there's any improvement there?

Author response:
**[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.

**As per Reviewer BjnM, we got new positive results with Markovian models (MBPO) in Appendix F and Figure 18 of the revised paper. In summary, DreamSmooth also benefits MBPO on the sparse-reward Hand task even though MBPO uses Markovian models. Together with the improvements of TD-MPC+DreamSmooth shown in Figure 9, this result clearly shows that DreamSmooth is generally applicable in model-based RL in practice.**

Thank you so much for your valuable feedback and we hope our new results address your concerns.

Reviewer comment (Evaluation):
There are inductive biases built into each smoothing function. For example, Gaussian and Uniform smoothing functions assume symmetry. However, this may not align with the environment dynamics. Consider a ball rolling off a staircase and receiving a sparse reward right at the edge of the staircase. The smoothing function bumps up the reward of the states before and after falling off the staircase, but in practice, it is much harder to climb back up from the lower platform than to roll down. In other words, reward smoothing can give rise to overoptimistic behaviors. Do you see this reflected in any task?

Author response:
**[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.

Reviewer comment (Experiments):
Does reward smoothing benefit model-free methods such as policy gradient?

Author response:
**[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.

Reviewer comment (Novelty):
**Novelty.** This is more like an engineering trick that the community has considered as an ad-hoc approach to resolve to learn in environments with sparse rewards. Although this smoothing technique intuitively makes sense, I didn’t see justifications for the correctness of this approach. See Question 1 below.

Author response:
**[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.

Reviewer comment (Experiments):
**Baseline method.** DreamV3 is the only baseline method for almost all the tasks, except that TD-MPC is used for the Hand Task. Unless this work only considers robotic tasks, other more popular RL algorithms need to be included. Also, if this work is indeed only constrained to robotic tasks, I believe the authors need to make that clear in the paper, and also explain why this simple technique cannot be applied to other RL algorithms.

Author response:
**[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.

Reviewer comment (Theory):
When we change the reward function, is the agent’s policy guaranteed to have a high value under the original reward function? When sparse rewards indeed specify critical states that have high rewards, would a smooth reward function blur out the true critical states, so that the optimal policy does not visit the critical states?

Author response:
**[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.

Thank you for your prompt response and we are happy to hear that our response addressed most of your concerns!

As the reviewer noticed, Theorem A.2 shows that our reward smoothing preserves optimal policies. However, a new MDP after reward smoothing cannot be defined in a conventional way since a smoothed reward function now additionally depends on a policy (i.e. future states and rewards). Therefore, we cannot provide a convergence guarantee to the optimal value function and policy in the smooth-reward setting, unlike in the original reward setting. 

Even so, our experiments indicate that this is not an issue in practice: DreamerV3 and TD-MPC successfully learn policies even with reward smoothing using future rewards, showing significant improvements on sparse-reward tasks while maintaining similar performance in dense reward environments.

Please let us know if you have any further questions or concerns!

Reviewer comment (Experiments):
Is there any rationale for using DreamV3 as the baseline for most tasks?

Author response:
**[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.