DreamSmooth: Improving Model-based Reinforcement Learning via Reward Smoothing
Authors: Vint Lee
Pieter Abbeel
Youngwoon Lee
Venue: ICLR 2024 poster (2024)
Keywords: Model-based Reinforcement Learning; Reward Shaping; Reward Smoothing
Primary area: reinforcement learning
TL;DR: We show that reward prediction is challenging in many sparse-reward tasks, and propose a simple yet effective method, reward smoothing, which effectively facilitates reward prediction and thus, improves model-based reinforcement learning.

Abstract: Model-based reinforcement learning (MBRL) has gained much attention for its ability to learn complex behaviors in a sample-efficient way: planning actions by generating imaginary trajectories with predicted rewards. Despite its success, we found that surprisingly, reward prediction is often a bottleneck of MBRL, especially for sparse rewards that are challenging (or even ambiguous) to predict. Motivated by the intuition that humans can learn from rough reward estimates, we propose a simple yet effective reward smoothing approach, DreamSmooth, which learns to predict a temporally-smoothed reward, instead of the exact reward at the given timestep. We empirically show that DreamSmooth achieves state-of-the-art performance on long-horizon sparse-reward tasks both in sample efficiency and final performance without losing performance on common benchmarks, such as Deepmind Control Suite and Atari benchmarks.

Published as a conference paper at ICLR 2024
DREAMSMOOTH: IMPROVING MODEL-BASED REIN-
FORCEMENT LEARNING VIA REWARD SMOOTHING
Vint Lee' Pieter Abbeel'
Youngwoon Lee YYnssii University
University of California, Berkeley
ABSTRACT
Model-based reinforcement learning (MBRL) has gained much attention for its ability to learn complex behaviors in a sample-effciient way: planning actions by generating imaginary trajectories with predicted rewards. Despite its success, we found that surprisingly, reward prediction is often bottleneck of MBRL, especially for sparse rewards that are challenging (or even ambiguous) to predict. Motivated by the intuition that humans learn from rough reward estimates, we propose a simple yet effective reward smoothing approach, DreamSmooth, which learns to predict semporally-smooteee given timestep. We empirically show that DreahSmooth achieves date-of-the-ar performance on long-horizon sparse-reward tasks both in sample efficiency and final performance without losing performance on common benchmarks, such as Deepmind Control Suite and Atari benchmarks.
reward, instead of the exact reward at the
INTRODUCTION
Humans often plan actions with rough
Tasks
Timestep
estimate of future rewards, instead the
exact reward at the exact moment (Fior-
Push block into bin
illo et al.) 2008)] Klein-Flügge et al.) 2011).
rough reward estimate is mostly suffi-
cient to learn a task, and predicting the ex-
act reward often challenging since it can be ambiguous, delayed, or not observable. Consider for instance the manipulation task (middle) of pushing
Target
illustrated in Figure | block on table into bin, where a large reward is given only on the timestep when block first touches the bin. Using the same image observations challenging even for humans to predict the correct sequence of rewards. Crucially, this issue is present in many environments, where states with no reward are almost in- distinguishable from those with rewards.
Rotate block
the agent, it
Dump rock into truck
No Reward
Large Reward
No Reward
Figure 1: Predicting the exact sequence of rewards is extremely difficult. These examples show the sequences of image observations seen by the agent just before and after it receives a large reward. There is little to visually distinguish timesteps with large reward from those without, which creates significant challenge for reward prediction.
An accurate reward model is vital model- based reinforcement learning (MBRL)- re- ward estimates that are too high will cause an agent to choose actions that perform poorly in reality, and estimates that are too low will lead an agent to ignore high re- wards. Despite its difficulty and importance, the reward prediction problem in MBRL has been largely overlooked. We find that even for the state-of-the-art MBRL algorithm. DreamerV3 (Hafner et al. 2023), reward prediction is not only challenging, but is also a performance bottleneck for many tasks. For instance, Dreamer V3 fails to predict any reward for most objectives in the Crafter environ- ment (Hafner 2022) with similar failure modes observed on variants of the RoboDesk (Kannan et al. 2021 and Shadow Hand (Plappert et al. 2018) tasks.

---

Published as a conference paper at ICLR 2024
Inspired by the human intuition that only rough estimate of rewards is sufficient, we propose simple yet effective solution, DreamSmooth, which learns to predict a semporally-smoottee rather than the exact reward at each timestep. This makes reward prediction much easier instead of having to predict rewards exactly, now the model only needs to produce an estimate of when large rewards are obtained, which is sufficient for policy learning.
reward
Our experiments demonstrate that while extremely simple, this technique significantly improves performance of different MBRL algorithms on many sparse-reward environments. Specifically, we find that for DreamerV3 (Hafner et al. 2023), TD-MPC (Hansen et al., 2022), and MBPO (Janner) al., 2019), our technique especially beneficial in environments with the following characteris- tics: sparse rewards, partial observability, and stochastic rewards. Finally, we show that even on benchmarks where reward prediction is not a significant issue, DreamSmooth does not degrade performance, which indicates that our technique can be universally applied.
2 RELATED WORK
Model-based reinforcement learning (MBRL) leverages a dynamics model (i.e. world model) of an environment and a reward model of desired task to plan a sequence of actions that maximize the total reward. The dynamics model predicts the future state of the environment after taking specific action and the reward model predicts the reward corresponding to the state-action transition. With the dynamics and reward models, an agent simulate large number of candidate behaviors in imagination instead of in the physical environment, allowing MBRL to tackle many challenging tasks (Silver al. 2016 2017 2018).
Instead of relying on the given dynamics and reward models, recent advances in MBRL have enabled learning a world model of dish-iiiennional observations and complex dynamics (Ha & Schmidhuber
2018, Schrittweee et al., 20200) Hafner et al. 2019 202 1 2023 Hansen et al., (022). as well as a
remporally-extente has achieved the state-of saate-of-the-ar the-art performance across diverse domains of problems, e.g., both with pixel and state observations as well as both with discrete and continuous actions.
world model (Shi et al.] 2022). Specifically, Dreamer V3 (Hafner et al. 2023)
For realistic imagination, MBRL. requires an accurate world model. There have been significant efforts in learning better world models by leveraging human videos (Mendonca al., 2023), by adopting more performant architecture (Deng et al. (023). and via representation learning, such as prototype-based (Deng et al. 2022) and object-centric (Singh et al. 202I) state representations, contrastive learning (Okada & Taniguchi| 2021), and masked auto-encoding (Seo et al. 2022 2023).
However, compared to the efforts on learning a better world model, learning an accurate reward model has been largely overlooked. [Babaeizadeh et al. (2020) investigates the effects of various world model designs and shows that reward prediction is strongly correlated to task performance when trained offline dataset, while limited to dense-reward environments. In this paper, point that accurate reward prediction is crucial for MBRL, especially in sparse-reward tasks and partially observable environments, and propose simple method to improve reward prediction in MBRL.
APPROACH
The main goal of this paper is understand how challenging reward prediction is in model-based reinforcement learning (MBRL) and propose simple yet effective solution, reward smoothing, which makes reward prediction easier learn. In this section, we first provide a background about MBRL in Section 3.1 then present experiments demonstrating the challenge of predicting sparse reward signals in Section 3.2 and finally explain our approach, DreamSmooth, in Section 3.4
3.1 BACKGROUND
We formulate problem partially observable Markov decision process (POMDP), which is defined as tuple (O, A, P. R, 7). O an observation space, A is an action space, P 1 agz) with timestep a transition dynamics, R is reward function that maps previous observations and actions to a reward r, = R(ogr, as)), and E [0, 1) is discount factor (Sutton & Barto 2018). RL
aims find policy π7 (a, oge; ax)) that maximizes the expected sum of rewards E2
mr).
2

---

higher loss than simply Published as a conference paper at ICLR 2024
This paper focuses MBRL algorithms that learn a world model Po (o (22411121, ae) and reward model Ro (re(zz) from agent experience, where &; is learned latent state at timestep t. The learned world model and reward model can then generate imaginary rollouts {zr, az, ]= 1+ of the horizon H starting from any 2t, which can be used for planning (Argenson & Dulac-Aroodd 2021] Hansen et al.) 2022) or policy optimization (Ha & Schmidhuber 2018 Hafner et al. 22019). Specifically, we use the state-of-the-art algorithms, DreamerV3 (Hafner al. 2023) and TD-MPC (Hansen et al. 2022)
D3eamerV3 (Hafner et al. 2023) uses the predicted rewards for computing new value targets to train the critic. For learning good policy, the reward model plays vital role since the critic. from which the actor learns policy, receives its training signal exclusively through the reward model. Note that the data collected from the environment only used for training a world model and reward model.
On the other hand, TD-MPC (Hansen et al., 2022) learns state-action value function Q(z az) directly from agent experience, not from predicted rewards. However, the reward model is still important for obtaining good policy in TD-MPC because the algorithm uses both the reward model and value function to obtain the policy through online planning.
3.2 REWARD PREDICTION IS DIFFICULT
Reward prediction is surprisingly challenging in many environments. Figure shows sequences of frames right before and after sparse rewards are received in diverse environments. Even for humans, it is difficult to determine the exact timestep when the reward is received in all three environments.
We hypothesize that the mean squared error loss (()g)) model training, deteriorates reward prediction accuracy when there exist sparse rewards. This is because predicting sparse reward single step earlier or later results predicting 0 reward at every step. Thus, instead of trying to predict sparse rewards at the exact timesteps, reward model minimizes the loss by entirely omitting sparse rewards from its predictions.
r) typically used for reward
To verify this hypothesis, we plot the ground-truth and DreamerV3's predicted rewards in Figure 2 On the four tasks described in Section 4.1 the reward models struggle at predicting exact rewards and simply ignore sparse rewards unless they are straightforward to predict. This hypothesis also holds deterministic and fully-observall rewards. The reward model fails to predict most of these reward sources (Figure 2d
environment, Crafter, which has 24 sources of sparse
The difficulty reward prediction can be further exacerbated by partial observability, ambiguous rewards, stochastic dynamics of environments. As an example in the first (third) in Figure
Ground Truth
Predicted
x-axis: Timestep
y-axis: Reward
0
-110
-110
120
120
(a) RoboDesk
(b) Hand
2400
-800
0.4
360
200
(c) Earthmoving
(d) Crafter
Figure 2: Ground truth rewards and DreamerV3's predicted rewards over an evaluation episode. The reward model misses many sparse rewards, which is highlighted in yellow.
3

---

Published as a conference paper at ICLR 2024
the sparse rewards given when the block (the rocks in the third example) first contacts the bin (the dumptruck). The exact moment of contact is not directly observable from the camera viewpoint, and this makes reward prediction ambiguous. Moreover, stochastic environment dynamics, e.g.. contact between multiple rocks, can make predicting a future state and reward challenging.
3.3 REWARD PREDICTION IS A BOTTLENECK OF MBRL
The preceding section shows that reward prediction is challenging many environments. More importantly, this poor reward prediction can be bottleneck of policy learning, as shown in Figure 3 In RoboDesk, where the reward model does not reliably detect the completion of the second task (Figure 2aj, the policy gets stuck at solving the first task and fails subsequent tasks. In Earthmoving, where the reward model cannot capture rewards for successful dumping (Figure (2c). the policy frequently drops the rocks outside the dumptruck. These consistent failure modes in reward prediction and policy learning in DreamerV3 suggest that poor reward prediction can be a bottleneck of MBRL.
Tasks Completed
Success
Rocks Sumpps
Stuck on 2nd task
Fallure
Finished task
16
24
8
12
Environment steps (x10f)
Environment steps (x109
(a) RoboDesk
(b) Earthmoving
Figure 3: The reward model's inability to predict sparse rewards for completing tasks leads poor task performance. (a) In RoboDesk, the agent gets stuck after learning the first task, and is unable to learn to perform subsequent tasks. (b) Earthmoving. policy often fails to dump the rocks accurately into the dumptruck. The learning curves are averaged over 3 seeds.
3.4 DREAM IMPROVING MBRL VIA REWARD SMOOTHING
To address the reward prediction problem, we propose a simple yet effective solution, DreamSmooth, which relaxes the require- ment for the model to predict sparse rewards at the exact timesteps by performing tempo- ral smoothing. Allowing the reward model to predict rewards that off from the ground truth by a few timesteps makes learning cas- ier, especially when rewards are ambiguous or sparse.
Dignnll a a
Orginl a-
5
Bennnd
mi
Timestop
Timestop
Timastop
(a) Gaussian
(b) Uniform
(c) EMA
Figure 4: Reward smoothing on sparse reward 1 at
t t= 4. 6, and smoothing hyperparameters.
Specifically, DreamSmooth applies temporal smoothing the rewards upon collecting each new episode. DreamSmooth can work with any smoothing function f that preserves the sum of rewards:
P, f(r Σ fs Tccip(t+i. s.t. E-=1,
(I)
where T and denote the episode and smoothing horizons, respectively. For simplicity, we omit the discount factor in Equation (I]: the full equation can be found in Appendix, Equation (6). Episodes with the smoothed rewards are stored the replay buffer and used to train the reward model. The agent learns only from the smoothed rewards, without ever seeing the original rewards. The smoothed rewards ease reward prediction by allowing the model to predict rewards several timesteps earlier or later, without incurring large losses. In this paper, investigate three popular smoothing functions: Gaussian, uniform, and exponential moving average (EMA) smoothing, as illustrated in Figure
While the main motivation for smoothing is to make it easier to learn reward models, note that reward smoothing in some cases preserves optimality an optimal policy under smoothed rewards also optimal under the original rewards r. particular, provide proof in Appendix A for the
4

---

Published as a conference paper at ICLR 2024
optimality of EMA smoothing (and any smoothing function where vi 0, 0) by augmenting the POMDP states with the history of past states. However, when future rewards are used for smoothing (e.g. Gaussian smoothing), the smoothed rewards conditioned policy, and we can longer define an equivalent POMDP. In such cases, there no theoretical guarantee. Even so, we empirically show that reward models can adapt their predictions alongside the changing policy, and achieve performance improvements.
The implementation of DreahSmooth is extremely simple, requiring only one additional line of code to existing MBRL algorithms, as shown in Algorithm [T] The overhead of reward smoothing is minimal, with time complexity O(T L). More implementation details can be found in Appendix B
Algorithm COLLECT MOLLLCCTTOOLLUUTT (æ:: policy, D: replay buffer) DREAMSMOOTH
{(o,, ai Fe) ROLLOUT(T) (re + -AUUSSAN((M) or (MA((r))E
D only one line needs to be added.
D DU ar, al,
4 EXPERIMENTS
In this paper, propose simple reward smoothing method, DreamSmooth, which facilitates reward prediction in model-based reinforcement learning (MBRL) and thus, improves the performance of existing MBRL methods. Through our experiments, we aim to answer the following questions: (1) Does reward smoothing improve reward prediction? (2) Does better reward prediction with reward smoothing lead to better sample efficiency and asymptotic performance of MBRL in sparse-reward tasks? (3) Does MBRL with reward smoothing also work in common dense-reward tasks?
4.1 TASKS
We evaluate DreamSmooth on four tasks with sparse subtask completion rewards and two common RL benchmarks. Earthmoving uses two 64 x images as an observation while all other tasks use a single image. See Appendix C for environment details.
RoboDesk: We use a modified version of RoboDesk (Kannan et al. 2021 where sequence of ma- push.green) need to be completed in order (Figure 5a). We use the original dense rewards together with large sparse reward for each task completed.
nipulation tasks (flat.block.ing
Hand: The Hand task (Plappert et al, 2018) requires Shadow Hand to rotate a block in hand into a specific orientation. We extend it to achieve sequence of pre-defined goal orientations order. In addition to the original dense rewards, we provide a large sparse reward for each goal.
Earthmoving: The agent controls a wheel loader to pick up rocks and dump them in the dump truck (Figure 5cp. Sparse rewards are given for picking up and dumping rocks, with dense rewards for moving rocks towards the dump truck. The environment is simulated using the AGX Dynamics physics engine (Algoryx 2020) with the AGX Terrain module (Servin et al. 202T).
Crafter: Crafter (Hafner 2022) is a minecraft-like 2D environment, where the agent tries to collect, place, and craft items in order to survive, There are 22 achievements in the environment
Complete in order
Dump
7 2
9
11
(a) RoboDesk
(b) Hand
(c) Earthmoving
(d) Crafter
(c) DMC
(f) Atari
Figure 5: We evaluate DreamSmooth on four tasks with sparse subtask completion rewards (a-d), We also test two popular benchmarks, (e) DeepMind Control Suite and (f) Atari.
5

---

Published as a conference paper at ICLR 2024
Ground Truth
Smoothed
Predicted
x-axis: Timestep
y-axis: Reward
0
0
-20
120
120
(a) RoboDesk
(b) Hand
200
0.2
360
200
(c) Earthmoving
(d) Crafter
Figure 6: We visualize the ground truth rewards, smoothed rewards with Gaussian smoothing, and predicted rewards by DreamerV3 trained the smoothed rewards an evaluation episode. In contrast to Figure the reward models with reward smoothing capture most of sparse rewards.
(e.g. collecting water, mining diamonds) with sparse reward for obtaining each achievement for the first time. A small reward is given (or lost) for each health point gained (or lost).
DMC:
We benchmark 7 DeepMind Control Suite continuous control tasks Tassa et al. 2018
Atari: We benchmark 6 Atari tasks (Bellemare et al. 2013) at 100K steps.
4.2
IMPROVED REWARD PREDICTION WITH REWARD SMOOTHING
We first visualize the ground truth rewards, smoothed rewards (Gaussian smoothing), and ward prediction results of DreamerV3 trained with DreamSmooth in Figure 6 We observe that ward smoothing leads to significant improve- ment in reward prediction: DreamSmooth success- fully predicts most of the (smoothed) sparse re- wards and no longer omits vital signals for policy learning or planning.
Baseline (no smoothing)
Gaussian
1.0
Prediction Rate
0.5
0.0
Health gained
Health loss
Collect cool
Collect drink
Collect sapling
Collect stone
Collect wood
Defeau Defeat skeleton eleton
Defeat zomble
Make stone picka claxe
Make stone sword
Eat cow
Make wood pic pickaxe
Make wob good sword
Place furnace
place plant
Place stone
Place table
Wake up
The improvement is especially notable in Crafter. In Figure 7 measure the accuracy of the re- ward model, (i.e. predicting reward larger than half of the original smoothed reward for Dream- crV3 and DreamSmooth respectively) at the exact timesteps for each subtask. The vanilla Dream- erV3's reward model (baseline) misses most of the sparse rewards while DreamSmooth predicts sparse rewards more accurately in 15/ 19 subtasks.
Figure 7: Reward prediction rates for 19 achieve- ments in Crafter. The other 3 tasks have been never achieved by both methods. With reward smoothing, the prediction rates are better in 15/19 19 tasks.
4.3 RESULTS
We compare the vanilla Dreamer (Hafner et al. 2023) with DreamSmooth, whose backbone is also DreamerV3. For DreamSmooth. evaluate Gaussian, uniform, and EMA smoothing. The hyperparameters for DreamerV3 and smoothing functions can found Appendix As shown in and CreamSmoott-nniforo significantly improve the performance
Figure

---

Published as a conference paper at ICLR 2024
DreamSmooth-
DreamSmooth-EMa
DreamerV3
Number of Rocks Dumped
Tasks Completed
Tasks Completed
2-
2-
1
Environment steps (x109)
Environment steps (x109)
Environment steps (x109)
(a) RoboDesk
(b) Hand
(c) Earthmoving
Human Normalized Score
12
(x10%)
0.6
0.2
Score
0.4
4-
0.2
0.0
0.0
1.00
0.3 Environment steps (x109)
Environment steps (x109)
Environment steps (x109)
(d) Crafter
(e) DMC
(l) Atari
Figure 8: Comparison of learning curves of DreamSmooth (Gaussian, Uniform, EMA) and Dream- erV3. The shaded regions in (a-d) show the maximum and minimum over 3 seeds. For DMC (e) and Atari (f), we aggregate results over 7 and tasks respectively, and display the standard deviation.
as well the sample efficiency of DreamerV3 on the Robodesk, Hand, and Earthmoving tasks. The only change between DreamerV3 and ours is the improved reward prediction, as shown in Section 4.2 This result suggests that reward prediction is one of major bottlenecks of the MBRL performance.
While all smoothing methods lead to improvements over DreamerV3, Gaussian smoothing generally performs the best, except Crafter, with uniform smoothing showing comparable performance. The better performance Gaussian and uniform smoothing could because allows predicting rewards both earlier and later, whereas EMA smoothing only allows predicting rewards later.
Despite the improved reward prediction accuracy, DreamSmooth- Saaammooth---aussian and DreamSmooth- Uniform perform worse than the baseline in Crafter. This could be because the symmetric Gaussian and Uniform smoothing kernels require the reward models to anticipate future rewards, while EMA smoothing does not. We believe this leads to more false-positive reward predictions from the former, leading to poor policy learning in Crafter. More details can be found in Appendix F
We also observe that on the DMC and Atari benchmarks, where reward prediction is not particularly challenging, our technique shows comparable performance with the unmodified algorithms (see Appendix, Figure 17 for full results), suggesting that reward smoothing can be applied generally.
In Figure g DreahSmooth also improves the performance of TD-MPC (Hansen et al. 2022) and MBPO (Janner et al. 2019). In the Hand task, the vanilla algo- rithms are unable to consistently solve the first task, even with proprioceptive state observations. However, DreamSmooth enables both algorithms to complete the tasks, even learning pixel observa- tions with TD-MPC. This suggests that DreamSmooth can be useful in broad range of MBRL algorithms that use a re- ward model. We only demonstrate the
Gaussian Smoothing Smoothing
x-axis: Eny steps (x10%) y-axis: Tasks Completed
0.0
0.0
(a) TD-MPC (Pixel) (b) TD-MPC (State) (c) MBPO (State)
Figure 9: Learning curves (median over 3 seeds) for TD- MPC and MBPO, with and without DreamSmooth, the Hand task.
7

---

Published as a conference paper at ICLR 2024
Hand task since TD-MPC and MBPO fail on other sparse- reward tasks, with MBPO requiring demonstrations to make progress on the Hand task (see Appendix for details).
4.4 ABLATION STUDIES
Data Imbalance. One possible cause of poor reward predictions is data imbalance - the reward model trains on few examples of sparse rewards due their infre- quency, potentially leading to poor predictions. To test this hypothesis. conducted experiments with over- sampling: with probability 0.5, we sample a sequence containing sparse rewards; otherwise, we sample uni- formly from all sequences in the buffer. As shown Figure TO, oversampling performs better than the baseline, but learns slower than DreamSmooth. This suggests that while data imbalance contributes to the difficulty of reward prediction, it is not the only fac- tor hindering performance. Furthermore, this over- sampling method requires domain knowledge about which reward signals are to be oversampled while DreamSmooth is agnostic to the scale and frequency of sparse rewards.
DreamSmooth Oversampling DreamerV3
Taskes Completed
Environment steps (x10")
Figure 10: Oversampling sparse rewards (p 0.5) improves DreamerV3 on Ro- boDesk, but still performs worse than Saannmmoot--Gassiiae The lines show median over 3 seeds, while shaded regions show maximum and minimum.
Reward Model Size. Another hypothesis for poor reward predictions is that the reward model does not have enough capacity to capture sparse rewards. To test this hypothesis, the size of the reward model from 4 layers of 768 units, which DreamSmooth uses, 5 layers of 1024 units and 6 layers of 1280 units, while keeping the rest of the world model the same. We observe in Figure T1 that without smoothing, increasing reward model size has negligible im- pact, and DreamSmooth outperforms all the ward model sizes tested. This indicates that the reward prediction problem is not simply caused by insufficient model capacity.
DreamSmooth DreamerV3-122
DreamerV3-1024 Dreamerr337768
Tasks Completed
1 Complee Pe
increase
12 18
24
20
Environment steps (x10%)
Environment steps (x109)
(a) RoboDesk
(b) Hand
Figure 11: Simply increasing the reward model size has negligible impact performance. DreamerV3-708 1024, and 1280 use 4, 5, layers of 768, 1024, 1280 units, respectively.
Loss Functions. We verify whether other formulations of reward regression could solve the reward prediction problem in RoboDesk. Following Hafner al. (2023), we take symlog the prediction target for stable training regardless of the scale of rewards. We find in Figure 12 that reward prediction remains challenge when using common loss functions, such as LI and L2, with LI significantly degrading task performance on RoboDesk. On the other hand, applying reward smoothing improves performance for all three loss functions.
Gaussian Smoothing
x-axis: steps (x109)
No Smoothing
y-axis: Tasks Completed
1
12
18
24
12
18
24
12
18
24
(a) 2-Hot
(b)
L1
(c) L22
Figure 12: Learning curves (median over seeds) of various loss functions for reward modeling on RoboDesk. Note that DreamerV3 and Dreamsmooth use the 2-Hot loss function.
analyze the impact of the smoothing parameters and a for Gaussian and EMA, respectively, on RoboDesk and Hand. We observe that DreamSmooth is insensitive to the smoothing parameters, performing well across wide range of values.
Smoothing Parameter. In Figure
8

---

Published as a conference paper at ICLR 2024
DreamSmooth Gaussian
0=
0
DreamerV3
Tasks Completed
Tasks Completed
2-
2
1
6
12
18
24
10
15
20
Environment steps (x109)
Environment steps (x10*)
(a) Gaussian Smoothing on RoboDesk
(b) Gaussian Smoothing on Hand
DreamSmooot-
DreamerV3
Tasks Completer
Tasks Completed
1.
12
Environment steps (x10%)
Environment steps (x109)
(c) Uniform Smoothing on RoboDesk
(d) Uniform Smoothing Hand
DreamSmooth-EMA
0.3
0.45
DreamerV3
Tasks Completed
Tasks Completed
0
12
18
24
Environment steps (x10%)
Environment steps (x10%)
(c) EMA Smoothing RoboDesk
(f) EMA Smoothing Hand
Figure 13: Parameter sweep over smoothing parameters a, ó, and a. The lines show median task performance over 3 seeds, while shaded regions show maximum and minimum.
CONCLUSION
In this paper, we identify the reward prediction problem MBRL and provide a simple yet effective solution: reward smoothing. Our approach, DreamSmooth, demonstrates superior performance in sparse reward tasks where reward prediction not trivial mainly due to the partial observability stochasticity of the environments. Moreover, DreamSmooth shows comparable results the commonly used benchmarks, DMC and Atari, showing its task-agnostic nature. Although we show that our simple reward smoothing approach mitigates the difficulty in reward prediction, the improved reward prediction does not always improve the task performance, e.g., in Crafter. This can be because more predicted task rewards can also result in more false positives. Further investigation on this trade-off is promising direction for future work.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was supported in part by the BAIR Industrial Consortium, an ONR DURIP grant, Komatsu, and InnoHK Centre for Logisties Robotics. We would like to thank Seohong Park for proofreading our proof and all members of the Berkeley Robot Learning lab for their insightful feedback.
REFERENCES
Algoryx. AGX dynamics, 2020. URL /tt::///wwwwwo httpppp/wwww.looor....///////e/e dynamics/
Arthur Argenson and Gabriel Dulac-Arnold. Model-based offline planning. In International Confer-
ence on Learning Representations, 2021. URL
forum?
OMNB1G5x2d4
Mohammad Babaeizadeh, Mohammad Taghi Saffar, Danijar Hafner, Harini Kannan, Chelsea Finn, Sergey Levine, and Dumitru Erhan. Models, pixels, and rewards: Evaluating design trade-off' in visual model-based reinforcement learning, arXiv preprint arXiv:2012.04603 2020.
M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, jun 2013.
Fei Deng, Ingook Jang, and Sungjin Ahn. Dreamerpro: Reconstruction-fre model-based reinforce- ment learning with prototypical representations. In International Conference on Machine Learning. pp. 4956-4975. PMLR, 2022.
Fei Deng, Junyeong Park, and Sungjin Ahn. Facing off world model backbones: Rnns, transformers. and s4. arXiv preprint urXiv:2307.02064 2023.
Christopher D Fiorillo, William T Newsome, and Wolfram Schultz. The temporal precision of reward prediction in dopamine neurons. Nature neuroscience, 11(8):96--73, 2008.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 10/22, 2018.
Danijar Hafner. Benchmarking the spectrum of agent capabilities. In International Conference on Learning Representations, 2022.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference Learning Representations, 2019.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference Learning Representations, 2021.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint 2rXiv:2301.0411 2023.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. In International Conference on Machine Learning, 2022.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. In Neural Information Processing Systems, 2019.
Harini Kannan, Danijar Hafner, Chelsea Finn, and Dumitru Erhan. Robodesk: A multi rein- forcement learning benchmark. 2021.
com/ cctt://ooooe--eeeerrcerhrbbbebbe
Miriam C Klein-Flicn KleineFlügge, Laurence T Hunt, Dominik R Bach, Raymond J Dolan, and Timothy EJ Behrens. Dissociable reward and timing signals in human midbrain and ventral striatum. Neuron, 72(4):654-664 2011.
Russell Mendonca. Shikhar Bahl, and Deepak Pathak. Structured world models from human videos. In Roboties: Science and Systems, 2023.
10

---

Published as a conference paper at ICLR 2024
Andrew Y Ng, Daishi Harada, and Stuart Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In International Conference on Machine Learning. volume 99, pp. 278-287, 1999.
Masashi Okada and Tadahiro Taniguchi. Dreaming: Model-based reinforcement learning by latent imagination without reconstruction. In IEEE International Conference on Robotics and Automation, pp. 1209-4215 2021. doi: 10.1 /....//CCRA48500200000.9.55
Matthias Plappert, Marein Andrychowiez, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, et al. Multi-goal reinforcement learn- ing: Challenging roboties environments and request for research. arXiv preprint (et::0904900000 2018.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604-699 2020.
Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Pieter Abbeel. Masked world models for visual control. In Conference on Robot Learning, 2022.
Younggyo Sco, Junsu Kim, Stephen James, Kimin Lee, Jinwoo Shin, and Picter Abbeel. Multi-view masked world models for visual robotic manipulation. In International Conference on Machine Learning, 2023.
Martin Servin, Tomas Berglund, and Samuel Nystedt. A multiscale model of terrain dynamics for real-time carthmoving simulation. Advanced Modeling and Simulation in Engineering Sciences, 8 (1):1-35, 2021.
Lucy Xiaoyang Shi, Joseph Lim, and Youngwoon Lee. Skill-based model-based reinforcement learning. In Conference on Robot Learning, 2022.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484
2016.
David Silver, Julian Schrittieserr Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of without human knowledge. nature, 350(7676):35- 359, 2017.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419): 1140-1144, 2018.
Gautam Singh, Skand Peri, Junghyun Kim, Hyunseok Kim, and Sungjin Ahn. Structured world belief for reinforcement learning in pomdp. In International Conference on Machine Learning, pp.
PMLR, 2021.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT Press, 2018
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmalcki, Josh Merel, Andrew Lefrancg, Timothy P. Lillicrap, and Martin A. Riedmiller. Deepmind control suite. arXiv preprint 1801 1006000 2018

---

Published as a conference paper at ICLR 2024
A PROOFS
Let M- (S be the given MDP. Without loss of generality, we assume the augmented form of the MDP M, where state 8, includes the entire history of states, i.e., 8; = (.81.....8)) 8t),
and thus, reward functions R, R have access to previous states, i.e. R((+t) = fi(s......... 8t).
Theorem A.I. An optimal policy #t of the MDP with reward smoothing only with past rewards, e.g.,
EMA smoothing, M = (S,A, P. 7) is also optimal under the original MDP M, where
R(s:) Σ fs f (((st)) and f-
(2)
(--h
Proof. We will use the theorem of reward shaping that guarantees optimal policy introduced in Ng ct al. (1999): if modified reward function can be represented in the form with any potential function (s:). the new reward function yields the same optimal policy with the original reward function R.
Let the potential function for the EMA reward smoothing
(8x)) = Σ R($++)) + Σ (((84)))
Σ fz-
(3)
1=
j=1+1
Then, our reward shaping term in R can be represented as the difference in the potential function (sc+1) (s:) as follows:
((8+))) (b(8)) -R(ss) Σ
Sa 7f R(8441).
1=-b
0
R(s:) -0000))+))--0000 = Σ
fir
R(8445) -R.
(5)
i=-L
Hence, following Ng et al. (1999), reward shaping with our EMA smoothing guarantees the optimal policy in the original MDP M.
However, Theorem A.I does not apply to smoothing functions that require access to future rewards, e.g., Gaussian smoothing. As in Gaussian smoothing, smoothed reward function may require future rewards, which are conditioned on the current policy; so the reward model. In such cases, there is no theoretical guarantee; but in experiments, empirically show that reward models can adapt their predictions along the changes in policies and thus, improve MBRL.
Instead, intuitively explain that optimal policy under any reward smoothing (even though the reward function is post hoc and cannot be defined for MDPs) also optimal under the original reward function.
Theorem A.2. An optimal policy म" with the smoothed reward function it is also optimal under the original reward function R. where
L Σ
R(s:) Σ 2 clp((..-.TT-)) R(s and
ft 11
(6)
12

---

Published as a conference paper at ICLR 2024
Proof. First, show that the discounted sum of original rewards Σi=o R(s)) and the of
smoothed rewards Li 7 R((s)) are the same for any trajectories (80,
T
2
(()))
= Σ
7
-lip(i,-t.T-) f: R(80li(((660077))
from Equation
6
(7)
1=0
1=
Σ
n¹ ((():))
Σ f
(8)
r R((s))
from
Σ
f; =
(9)
t=
Let an optimal policy under the smoothed rewards it be π* Assume that π is not optimal under the original reward R. Then,
T
3m 80
such that
E 3r)xx"
R(se)
>
80 or)
R((s))
(10)
1=
However,
T
E 60 ST) "
Σ
7
[((s))
13 (T)
7 (((s))
by Equation 9
(11)
1=0
t=
E (o......)
72 R(s)))
by Equation 110
(12)
which contradicts that n* is optimal under n.. Therefore, the optimal policy #* under ň guarantees its optimality under R.
B IMPLEMENTATION DETAILS
Models are trained on NVIDIA A5000, V100, RTX Titan, RTX 2080, and RTX 6000 GPUs. Each experiment takes about 72 hours for RoboDesk, 100 hours for Hand, 150 hours for Earthmoving, 96 hours for Crafter, and 6 hours for Atari and DMC tasks.
B.1 SMOOTHING FUNCTIONS IN DREAMSMOOTH
Gaussian smoothing follows the Gaussian distribution with a::
fi ke 20
(13)
where k = 1/(E - normalization constant.
We implement this using
scipy
aalilannffffrrdd (rewards, sigma, Tode="nearest"
Uniform smoothing distributes rewards equally across N consecutive timesteps.
5-1 5-1
fo === Vi E
(14)
2 2
We implement this using
scipy- indimage.... convol lonvovee veeeleeereewwaa (rewards, filter, mode="nearest:
EMA smoothing uses the following smoothing function:
f = - a) Vi S
(15)
which we implement by performing the following at each timestep:
reward[t] al 1pha reward(t
1] (1 aalhaa reward
13

---

Published as a conference paper at ICLR 2024
B.2 MODEL-BASED REINFORCEMENT LEARNING BACKBONES
Hyperparameters for DreamerV3, TD-MPC, and MBPO experiments are shown in Table 1 Table 2
and T respectively.
Table 1: DreamerV3 hyperparameters. Episode length is measured in environment steps, which is the number of agent steps multiplied by action repeat. Model sizes are as listed in Hafner et al. (0023). which we also refer for all other hyperparameters.
Environment
Action Repeat
Episode Length
Train Ratio
Model Size
5
Earthmoving
2000
64
3
0.33
RoboDesk
8
2400
64
L
3
0.3
9
Hand
300
64
3
0.3
Crafter
1
Variable
64
XL
3
0.3
9
DMC
2
1000
512
3
0.33
Atari
4
Variable
1024
S
3
0.3
9
Table 2: TD-MPC hyperparameters. We refer to Hansen et al.
2022
for all other hyperparameters.
Environment
Latent Dimension
CNN channels
Planning Iterations
Hand-Pixo
128
64
Hland-Proprio
128
12
Table MBPO hyperparameters. We refer to Janner al. (219 for all other hyperparameters.
Environment
Layer Size
Prediction Head Size
Demo Pre-Traine Pre-Training Steps
Hann-Proprio
512
400
30000
RoboDesk-Prope Proprio
512
400
30000
3
DMC-Proprio
256
200
0
3
C ENVIRONMENT DETAILS
C.I ROBODESK ENVIRONMENT
We use modified version of RoboDesk (Knnan et al. 202 where sequence of manipulation tasks fclatbblock in bin, upr apright.ar block..ck..ooo bww.kkoof.taaaee.e..e off.table.............o.. push-green need to be completed in order. Figure 14 shows images of an agent successfully completing each of these tasks.
In the original environment, dense rewards are based on Euclidean distances of objects to their targets, with additional terms encourage the to reach the object. They typically range from 0 to 10 per timestep. We use these dense rewards together with large sparse reward of 300 for each task completed.
C.2 HAND ENVIRONMENT
We modified the Shadow Hand environment (Plappert et al. 2018), so that the agent is required to achieve sequence of pre-defined goal orientations in order. The first goals are shown in Figure [5] while the subsequent goals agent only has to rotate the cube along the z-axis, and we only require the agent to match the cube's 's rotation the goal, not its position.
repeat of the first 3. The goal orientations are chosen that the
In the original environment, dense rewards are computed using r = - -10x + Ad), where r the Euclidean distance to some fixed position, and A is the angular difference the target orientation. In addition to these dense rewards, we provide a large sparse reward of 300 for each goal successfully achieved by the agent.
14

---

Published as a conference paper at ICLR 2024
Push green block into the bin
(b) Push teal block off the table
(c) Press the green button
Figure 14: Subtasks for RoboDesk.
I
C
I
(a) Goal 1
(b) Goal 2
(c) Goal
Figure 15: Subtasks for Hand.
C.3 AGX EARTHMOVING ENVIRONMENT
The Earthmoving environment consists of wheel loader, dump truck, a pile of dirt, with some rocks on top of the pile. The en- vironment is simulated using the realistic AGX Dynamics physics engine (Algoryx) 2020). The agent controls the wheel loader to pick up rocks and dump them in the dump truck.
The starting positions of the dirt pile, wheel loader, and dump truck the initial orientations of the dirt pile and
all randomized, wheel loader.
The agent uses one camera mounted on the cabin (left) for navigation, and one mounted on the bucket (right) for observing interactions with rocks and terrain.
Figure 16:
The agent's observations consist of components: wide-angle egocentric RGB camera mounted on the cabin to allow navigation, an RGB camera mounted on the bucket for observing interactions with rocks, and proprioceptive observations (positions, velocity, speed, force of actuators etc.). use 64 x 3 images for all cameras, while the proprioceptive observation has 21 dimensions.
The action space is 4-dimensional: dimensions for driving and steering the loader, and dimensions for moving and tilting the bucket.
The reward consists of a large sparse reward for rocks picked and dumped, and dense rewards for moving rocks towards the dumptruck. The total reward o at timestep / is computed using Equation [6).
Adump (maump
m 1-1 + Alead (mo Mionnd
Amove fsod (max (2, d') max (2, d
sparse reward
dense reward
(16)
Where Mauump Mead are rock masses the dumptruck and the bucket respectively, is the distance between the shovel and a point above the dumptruck, and À constants.
15

---

Published as a conference paper at ICLR 2024
D DMC AND ATARI BENCHMARKING RESULTS
DreamSmooth-Gaussian
DreamSmmotth-nnffor
DreamSmooth-EMe
DreamerV3
Hopper Hop
Reacher Hard
Cartpole Swingup Sparse
1000
Score
Score
Score
500
0.0
0.5
Environment steps (x10%)
Environment steps (x10%)
Environment steps (x10*)
Walker Run
Finger Turn Hard
Quadruped Run
1000
Score
Score
Score
500
500
500
0-
0.5
1.0
Environment steps (x10%)
Environment steps (x10%)
Environment steps (x10")
Cheetah Run
Score
500
0.0
0.5
Environment steps (x10%)
Pong
Breakout
Freeway
-9.00
13:00
Score
-13.33
Score
8-
Score
8.33
-17.67
3.67
-22.00
-1.00
0.2
Environment steps (x10*)
Environment steps (x10*)
Environment steps (x10")
Assault
Seaquest
Hero
601
339
8692
Score
489
243
Score
5862
377
147
3033
265
51
203
0.0
Environment steps (x10*)
Environment steps (x10%)
Environment steps (x109)
Figure 17: Full learning curves for the DMC and Atari benchmarks.
16

---

Published as a conference paper at ICLR 2024
E MBPO
In addition to DreamerV3 and TD-MPC, we show the general applicability of our reward smoothing method with other MBRL algorithms, such as Model-Based Policy Optimization (MBPO) (Janner et al. 2019).
Similar to the experiments on DreamerV3, Figure [ demonstratess that reward smoothing does affect the performance of MBPO on the dense-reward DMC tasks (Walker Walk, Cheetah Run, Reacher Easy. Cartpole Swingup). while significantly improving the performance of MBPO on sparse-reward tasks, especially in the Hand environment. Since MBPO, when trained from scratch, struggles at solving Hand and Robolesk, we initialize the replay buffer for these two experiments using trajectories from the fully-trained DreamerV3 policies (100 and 68 episodes, respectively). MBPO is then able to learn the first task Hand with reward smoothing, while failing at both Hand and RoboDesk without reward smoothing. This indicates that DreamSmooth can also benefit Markovian models that have no access to neither past nor future states.
TreamSmmotth-aussiian (MBPO)
MBPO
Walker Walk
Cheetah Run
Reacher Easy
1000
Score
Score
Score
500
500
500
0.00
0.10
0.00
0.15
0.05
0.15
Environment steps (x10%)
Environment steps (x10%)
Environment steps (x109)
Cartpole Swingup
Hand
Robodesk
1000
Tasks Completed
1.0
Tasks Completed
1.0
Score
500
0.5
0.5
0
0.0
AAAA
0.10 Environment steps (x10%)
0.04
0.00
0.15
Environment steps (x10*)
Environment steps (x10')
Figure 18: Learning curves for MBPO with DreamSmooth (Gaussian smoothing) on Hand, RoboDesk. and DMC tasks. For Hand and RoboDesk, MBPO trained from scratch could not achieve any meaningful reward signal. To ease the exploration problem, we initialize the replay buffers with demonstrations collected from the fully-trained DreamerV3 policies. The shaded regions show the minimum and maximum over seeds.
17

---

Published as a conference paper at ICLR 2024
F CRAFTER
is surprising that despite improved reward predictions when using Gaussian or uniform smoothing on the Crafter environment, the task performance significantly deteriorates compared to vanilla DreamerV3, or even DreamerV3 with EMA reward smoothing, as can be seen in Figure 8d
One possible cause is the symmetrical structure of the Gaussian and Uniform kernels, which makes the smoothed rewards dependent on both past and future ground truth rewards. This means that the reward model has to anticipate future rewards when performing predictions. We suspect that this leads to high rate of false positives in Crafter, where there are many sources of sparse rewards. The false positives can make policy learning difficult.
To test this, we introduce 2 asymmetric variants of the uniform smoothing kerne::
Uniform [-5, 0], where the smoothed rewards only depend on past ground truth rewards.
1
Vi [-
(17)
0+1
Uniform [0, 8). where the smoothed rewards only depend on future ground truth rewards.
1 fi 8+1 Vi [0,8].
(18)
The plots of predicted and ground truth rewards in Figure 20 and Figure 2T show that the reward models trained with the symmetric uniform smoothing kernel and Uniform [0, 4) smoothing kernel have more false positives than those trained with Uniform [-4, 0] and no smoothing, predicting significant positive rewards even when ground truth reward is 0.
In Figure 19 smoothing kernels that require the model to anticipate future rewards perform worse than those that do not: Uniform |-4, 0] performs the best, followed by symmetric uniform smoothing. then Uniform [0, 4].
Uniform (Symmetric)
Uniform [-4, 0]
Uniform [0, 4]
No Smoothing
16
Score
0.0
1.5
3.0
4.5
Environment steps (x10*)
Figure 19: Learning curves of DreamSmooth with uniform smoothing variants on Crafter. Perfor- mance on Crafter highly correlated with the need to anticipate future reward.
18

---

Published as a conference paper at ICLR 2024
Ground Truth
Smoothed
Predicted
x-axis: Timestep
y-axis: Reward
0.6
1.2
0.0
0.0
-0.2
-0.4
0
200
200
0.0
-0.2
-0.4
200
400
0.6
1.2
0.0
0.6
1.2
0.0
0.0
-0.2
-0.4
0
100
200
0.0
-.
-
100
300
(a) Uniform [-4, (Symmetric) Smoothing
(b) Reward Smoothing
Figure 20: Ground truth rewards and predicted rewards over five evaluation episodes Crafter. Similar to Uniform [0, 4] smoothing, the reward model trained with symmetric uniform smoothing predicts many false positives. Mcanwhile. the reward model trained without smoothing predicts few false positives but has many false negatives.
19

---

Published as a conference paper at ICLR 2024
Ground Truth
Smoathed
Predicted
x-axis: Timestep
y-axis: Reward
0.0
0.2
-0.2
400
100
0.6
0.6
-0.
0.6
0.6
0.0
0.0
0.2
0.2
400
100
0.6
0.0
-02
-0.2
200
100
0.6
-0.2
(a) Uniform [-4, 0) Smoothing
(b) Uniform |0, 4] Smoothing
Figure 21: Ground truth rewards and predicted rewards over five evaluation episodes on Crafter. With Uniform O0, 4] smoothing, the reward model has anticipate future rewards, resulting in more false positives than Uniform |44 0), where the smoothed rewards at the current timestep only depend on past rewards.
20


Reviews:
[{"review_id": "aN8uYm4f6f", "reviewer": "Reviewer_4wgT", "summary": "In model-based reinforcement learning (MBRL), it is crucial to correctly estimate the reward model. However, when the rewards in the environment are sparse, it poses challenges in learning the reward function. The authors have shown convincing examples that the algorithm may achieve a smaller loss by simply predicting zero rewards, than predicing the sparse rewards at an incorrect time step.\n\nThis work remedies this problem by asking the algorithm to learn a smooth reward function. The proposed algorithm is based on the DreamerV3 algorithm and evaluated on a wide range of tasks.", "questions": "1. When we change the reward function, is the agent’s policy guaranteed to have a high value under the original reward function?\nWhen 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?\n\n2. Is there any rationale for using DreamV3 as the baseline for most tasks?\n\n---\n\nMy questions and concerns about weaknesses are addressed in the rebuttal.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This work proposes a simple yet effective approach to improve an MBRL agent’s performance in environments with sparse rewards. The algorithm is evaluated on simulated robotic control, 2D navigation, and Atari game domains. The authors also conducted ablation studies to show that this approach outperforms some other baseline algorithms to address the sparse reward issue, including oversampling sequences with sparse rewards, increasing reward model size, etc.\n\nAdditionally, this work also empirically verifies the challenges of reward learning in MBRL, finding out that the agent may achieve a smaller loss by predicting zero rewards than by predicting wrong rewards.", "weaknesses": "**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.\n\n**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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "buQga83YGt", "reviewer": "Reviewer_BjnM", "summary": "This paper introduces DreamSmooth, a simple and effective method that improves the performance of model-based RL on sparse reward environments. The authors observe that on sparse reward tasks, it is challenging to fit an accurate reward model due to data imbalance. This in turn bottlenecks the performance of model-based RL methods like Dreamer. To mitigate this problem, the authors apply a smoothing function to the reward, effectively spreading the reward signal to adjacent states in the trajectory. The authors propose three reward smoothing functions: Gaussian, Uniform, and EMA. While EMA is the only function that guarantees policy invariance, all three are empirically found to work well, resulting in more accurate reward predictions and higher performance in sparse reward tasks compared to baselines.", "questions": "- 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?\n- 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?\n- Does reward smoothing benefit model-free methods such as policy gradient?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Reward sparsity is a long-standing problem in model-based RL. Even with an accurate dynamics model, policy optimization would not work if the reward model fails to capture sparse reward signals.\n- The method is extremely simple, with only a one-line change to the Dreamer code, yet it brings a significant improvement across a suite of challenging sparse reward tasks.\n- The authors performed extensive analysis and ablation studies to identify the root cause of MBRL failure and demonstrate the effectiveness of their algorithm.", "weaknesses": "- 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.\n- Two of the smoothing functions are unable to guarantee policy invariance, and it is possible to construct adversarial examples (see questions).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uJmDhVokRK", "reviewer": "Reviewer_7VX8", "summary": "The paper studies the challenges of learning reward models in the context of Model-Based Reinforcement Learning (MBRL). The authors argue that existing methods in MBRL fail to learn a good reward model in sparse / partially observable/stochastic environments/tasks and show empirical evidence for this. Based on the intuitive idea that, in such challenging scenarios, one only has to rely on rough reward estimates as humans do, they propose a method called DreamSmooth. In DreamSmooth the reward model is now tasked to predict a temporally smoothed reward instead of the exact reward. The authors propose 3 reward-smoothing schemes and empirically show that the approach can significantly improve the performance on most sparse reward scenarios.", "questions": "1) 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?\n2) The rationale behind why the counter-intuitive results in crafter is not convincing. Did the authors perform further empirical studies / analysis ?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "These are the strengths of the paper in my opinion:\n\n1) Studies a largely ignored, yet important problem of reward modelling in the context of MBRL.\n2) Propose a simple yet effective solution for the same.\n3) Well Written.", "weaknesses": "The major weaknesses are as follows:\n\n1) Counter-intuitive results in Crafter, where the method performs worse even after having a much better reward model.\n2) 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aRsZs2ihMs", "reviewer": "Reviewer_Q6F6", "summary": "This paper propose DreamSmooth which is based on Dreamer-V3. The technique used in this paper is to smooth the ground truth rewards when trainining Dreamer-V3. Experiments show that this modification works well on both dense and sparse rewards environments.", "questions": "I am curious about why DreamSmooth works comparable without reward smooth technique in dense reward environments. The reward smoothing technique changes the  reward distribution. As far as I am concerned, using this kind of technique would decline the final performance. Since MBRL like Dreamer is so important to this community, I lean to weak accept for this paper. However, I do have some concerns about this simple yet effective algorithms, especially why the final performance doesn't decline (especially on dense reward environments).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The proposed method solves a problem in model based RL when the gt rewards are sparse in these environments. Smoothing the rewards would make the reward function prediction process (reward leanring in model based RL) much better than before. At the same time, DreamSmooth shows it also performs well on dense rewards environments. It makes the whole algorithm more convincing. I think it is a good paper to investigate the reward smoothing technique for model based RL.", "weaknesses": "Dreamer-V3 has a symlog prediction function with reward learning process. I think different reward prediction function would contribute to the reward learning process. Do the authors conduct some experiments with different prediction function head to justify whether it could solve the sparse rewards problem? I think the issue discussed in this paper is mainly about the reward generalizability problem. It is hard for reward function in MBRL to generalize in sparse reward setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)