PaperlayerView exact revision on arXiv

Source-linked AI summary

DreamSmooth: Improving Model-based Reinforcement Learning via Reward Smoothing

Vint Lee, Pieter Abbeel, Youngwoon Lee

arXiv:2311.01450v2cs.LGcs.AIcs.RO

TL;DR

MBRL can fail when sparse rewards are difficult to predict accurately, making reward modeling a bottleneck. DreamSmooth replaces exact per-timestep reward prediction with temporal smoothing, and improves sparse-task performance while maintaining comparable results on common benchmarks. Future-looking smoothing has no theoretical optimality guarantee, and improved reward prediction can still produce false positives in some settings.

  • Problem

    Sparse and ambiguous rewards make accurate reward prediction difficult, yet reward modeling is important for MBRL performance.

    Source: §1 · p. 1, §2 · p. 2

  • Method

    DreamSmooth temporally smooths rewards before training the reward model, allowing predictions to occur near rather than exactly at reward timesteps.

    Source: §1 · p. 2, §3.4 · p. 4

  • Results

    DreamSmooth significantly improves multiple MBRL algorithms on sparse-reward tasks while maintaining comparable performance on DMC and Atari benchmarks.

    Source: §1 · p. 2, §4.3 · p. 7, §5 · p. 9

  • Takeaways & Limitations

    Reward smoothing is a simple, broadly applicable way to mitigate difficult reward prediction in sparse, partially observable, or stochastic environments.

    Source: §1 · p. 2, §5 · p. 9

    Future-reward smoothing lacks a theoretical optimality guarantee, and improved reward prediction can sometimes produce false positives that harm task performance.

    Source: §3.4 · p. 5, §4.3 · p. 7

Abstract

from arXiv · show

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.

1 INTRODUCTION

Reward prediction is a major MBRL challenge because sparse rewards can be ambiguous and visually indistinguishable across timesteps. DreamSmooth addresses this by predicting temporally smoothed rewards, improving sparse-task performance without degrading common benchmarks.

1 INTRODUCTION

  • Sparse rewards are difficult to predict because observations at rewarded and unrewarded timesteps can look nearly identical.The problem is especially challenging when rewards are delayed, ambiguous, or not directly observable.Source: §1 · p. 1, Figure 1 · p. 1
  • Reward prediction is a performance bottleneck in DreamerV3 and related tasks, where inaccurate estimates can cause poor action selection or ignored rewards.DreamerV3 fails to predict many objectives in Crafter, with similar failure modes in RoboDesk and Shadow Hand.Source: §1 · p. 1
  • DreamSmooth learns temporally smoothed rewards, allowing the model to estimate when large rewards occur instead of matching their exact timesteps.This relaxes exact reward-timing requirements while retaining a signal sufficient for policy learning.Source: §1 · p. 2
  • DreamSmooth significantly improves multiple MBRL algorithms on sparse-reward environments, especially with sparse, partially observable, or stochastic rewards.The evaluated algorithms are DreamerV3 (Hafner et al., 2023), TD-MPC (Hansen et al., 2022), and MBPO (Janner et al., 2019).Source: §1 · p. 2
  • On benchmarks where reward prediction is not difficult, DreamSmooth does not degrade performance, supporting application across common MBRL benchmarks.The paper reports comparable performance on DeepMind Control Suite and Atari benchmarks.Source: §1 · p. 2

2 RELATED WORK

MBRL combines learned world and reward models to plan through imagined trajectories, while prior work has emphasized world-model quality more than reward prediction. This paper highlights reward modeling as especially important in sparse-reward and partially observable settings.

2 RELATED WORK

  • MBRL uses dynamics and reward models to simulate candidate behaviors in imagination and plan actions that maximize predicted cumulative reward.The dynamics model predicts future states, while the reward model evaluates state-action transitions.Source: §2 · p. 2
  • Recent MBRL advances improve world models for high-dimensional observations and complex dynamics, with DreamerV3 achieving state-of-the-art performance across diverse domains.These domains include pixel and state observations and discrete and continuous actions.Source: §2 · p. 2
  • World-model research has explored human videos, architectures, representation learning, contrastive learning, and masked auto-encoding to improve realistic imagination.These approaches target better representations or predictive modeling rather than the reward model itself.Source: §2 · p. 2
  • Reward-model accuracy has received less attention, although prior evidence links reward prediction to task performance mainly in dense-reward offline settings.This paper extends the focus to sparse-reward and partially observable environments and proposes reward smoothing.Source: §2 · p. 2

3 APPROACH

The approach identifies sparse-reward prediction as a source of MBRL failure and introduces DreamSmooth, which temporally smooths rewards before training the reward model. The method relaxes exact timing requirements while raising theoretical caveats for future-looking smoothing.

3.1 BACKGROUND

3.2 REWARD PREDICTION IS DIFFICULT

  • Exact sparse-reward prediction is difficult because rewarded and unrewarded observations can be visually indistinguishable, especially under partial observability or stochastic dynamics.Contact events in RoboDesk and Earthmoving may be ambiguous from camera observations, while stochastic interactions further complicate prediction.Source: §3.2 · pp. 3–4, Figure 2 · p. 3
  • Mean squared error can encourage reward models to omit sparse rewards, and DreamerV3 misses many reward sources across evaluated tasks.A prediction shifted by one timestep can incur greater loss than predicting zero throughout, favoring omission.Source: §3.2 · p. 3

3.3 REWARD PREDICTION IS A BOTTLENECK OF MBRL

  • Poor reward prediction can bottleneck policy learning: DreamerV3 gets stuck on early RoboDesk tasks and often misses successful dumping in Earthmoving.These linked reward-model and policy failures motivate treating reward prediction as a central MBRL problem.Source: §3.3 · p. 4, Figure 3 · p. 4

3.4 DREAMSMOOTH: IMPROVING MBRL VIA REWARD SMOOTHING

  • DreamSmooth temporally smooths rewards so predictions may occur a few timesteps early or late, easing learning for sparse or ambiguous rewards.The method applies smoothing when collecting episodes, stores only smoothed rewards, and evaluates Gaussian, uniform, and EMA functions.Source: §3.4 · p. 4, Figure 4 · p. 4
  • Smoothing functions preserve total reward, while implementation adds one code line with O(T · L) time complexity.T and L denote the episode and smoothing horizons.Source: §3.4 · pp. 4–5
  • EMA smoothing has an optimality proof under an augmented-state formulation, but future-reward smoothing such as Gaussian smoothing lacks an equivalent POMDP guarantee.Future-looking smoothing conditions rewards on the policy, although empirical performance can still improve.Source: §3.4 · pp. 4–5

4 EXPERIMENTS

The experiments test DreamSmooth across sparse-reward tasks and standard benchmarks, showing improved reward prediction and stronger MBRL performance while exposing smoothing-specific limitations and ablation findings.

4.1 TASKS

  • Experiments evaluate DreamSmooth on four sparse subtask-completion environments—RoboDesk, Hand, Earthmoving, and Crafter—and benchmark 7 DeepMind Control Suite tasks and 6 Atari tasks.Earthmoving uses two 64 × 64 images as an observation, while all other tasks use a single image.Source: §4.1 · pp. 5–6

4.2 IMPROVED REWARD PREDICTION WITH REWARD SMOOTHING

  • DreamSmooth improves reward prediction in sparse-reward settings, predicting rewards more accurately in 15/19 Crafter subtasks.The smoothed reward model captures most sparse rewards and avoids omitting signals useful for policy learning or planning.Source: §4.2 · p. 6, Figure 7 · p. 6

4.3 RESULTS

  • DreamSmooth improves DreamerV3 performance and sample efficiency on RoboDesk, Hand, and Earthmoving, while Gaussian smoothing generally performs best.Uniform smoothing is comparable, whereas Crafter is an important exception where Gaussian and Uniform underperform the baseline.Source: §4.3 · p. 7
  • Crafter limits the benefit of symmetric smoothing: Gaussian and Uniform improve reward prediction but perform worse than vanilla DreamerV3, likely because false-positive predictions arise from anticipating future rewards.EMA does not show the same reported deterioration in Crafter.Source: §4.3 · p. 7
  • On DMC and Atari, where reward prediction is less challenging, DreamSmooth has comparable performance to the unmodified algorithms.The evaluation aggregates seven DMC tasks and six Atari tasks.Source: §4.3 · p. 7, Figure 8 · p. 7
  • DreamSmooth also improves TD-MPC and MBPO on the Hand task, enabling both algorithms to complete tasks that vanilla versions cannot consistently solve.TD-MPC learns from pixel observations with DreamSmooth, but the comparison is demonstrated only on Hand because the baselines fail on other sparse-reward tasks.Source: §4.3 · pp. 7–8, Figure 9 · p. 7

4.4 ABLATION STUDIES

  • Oversampling sparse-reward sequences improves over baseline but learns more slowly than DreamSmooth and requires domain knowledge about which signals to oversample.This indicates data imbalance contributes to, but does not fully explain, reward-prediction difficulty.Source: §4.4 · p. 8, Figure 10 · p. 8
  • Increasing reward-model capacity has negligible impact without smoothing, while DreamSmooth outperforms the tested larger models and remains effective across a wide parameter range.The parameter sweep finds DreamSmooth insensitive to smoothing parameters across RoboDesk and Hand.Source: §4.4 · p. 8, Figure 11 · p. 8, Figure 13 · p. 9

5 CONCLUSION

DreamSmooth identifies reward prediction as an MBRL problem and uses reward smoothing to improve sparse-reward performance while remaining comparable on DMC and Atari.

  • DreamSmooth improves sparse-reward performance, especially when reward prediction is difficult because environments are partially observable or stochastic.Source: §5 · p. 9
  • DreamSmooth remains comparable on commonly used DMC and Atari benchmarks, indicating task-agnostic behavior.Source: §5 · p. 9
  • Improved reward prediction does not always improve task performance, as shown by false-positive rewards in Crafter.Source: §5 · p. 9

A PROOFS

The proofs establish optimality preservation for past-only reward smoothing, while future-reward smoothing lacks a theoretical guarantee because it depends on policy-conditioned rewards.

  • EMA smoothing preserves optimal policies under the original reward when the MDP state is augmented with past-state history.Source: A PROOFS · p. 12
  • The proof represents EMA smoothing as potential-based reward shaping, invoking the condition R(s_t)+γΦ(s_{t+1})−Φ(s_t).Source: A PROOFS · p. 12
  • Future-reward smoothing, such as Gaussian smoothing, has no theoretical guarantee because smoothed rewards are conditioned on the current policy.Source: A PROOFS · p. 12
  • The proof argues that discounted original and smoothed returns coincide for trajectories, so an optimum under smoothed rewards is also optimal under original rewards.Source: A PROOFS · p. 13

B IMPLEMENTATION DETAILS

DreamSmooth implementation uses simple smoothing functions and trains the reward model on episodes containing smoothed rewards, with minimal computational overhead.

B IMPLEMENTATION DETAILS

  • DreamSmooth adds minimal overhead, requiring about 72–150 hours for the listed environments and 6 hours for Atari and DMC experiments.Source: B IMPLEMENTATION DETAILS · p. 13

B.1 SMOOTHING FUNCTIONS IN DREAMSMOOTH

  • Gaussian smoothing follows a Gaussian distribution, while uniform smoothing distributes rewards equally across δ consecutive timesteps.Source: §B.1 · p. 13
  • EMA smoothing updates each reward as reward[t] = alpha * reward[t - 1] + (1 - alpha) * reward[t].Source: §B.1 · p. 13

B.2 MODEL-BASED REINFORCEMENT LEARNING BACKBONES

The experiments use DreamerV3, TD-MPC, and MBPO across sequential sparse-reward manipulation and earthmoving tasks, with image-based observations and mixed sparse-dense rewards.

B.2 MODEL-BASED REINFORCEMENT LEARNING BACKBONES

C.1 ROBODESK ENVIRONMENT

  • RoboDesk requires completing flat-block, upright-block, and push-green manipulation tasks sequentially, combining dense rewards with a sparse reward of 300 per task.Source: §C.1 · p. 14

C.2 HAND ENVIRONMENT

C.3 AGX EARTHMOVING ENVIRONMENT

  • Earthmoving uses a randomized wheel-loader environment in which the agent picks up rocks and dumps them into a truck using two cameras and proprioceptive observations.Source: §C.3 · p. 15, Figure 16 · p. 15
  • Earthmoving combines sparse rewards for picking up and dumping rocks with dense rewards for moving rocks toward the dump truck.Source: §C.3 · p. 15

D DMC AND ATARI BENCHMARKING RESULTS

Figure 17 presents the full learning curves for the DMC and Atari benchmarks.

  • Figure 17 reports full learning curves for the DMC and Atari benchmarks.Source: Figure 17 · p. 16

E MBPO

DreamSmooth extends reward smoothing to MBPO, preserving dense-task performance while improving learning on sparse-reward tasks. On Hand, it enables learning where MBPO without smoothing fails.

E MBPO

  • The MBPO experiments demonstrate DreamSmooth’s applicability beyond DreamerV3 and TD-MPC.Source: E MBPO · p. 17
  • DreamSmooth improves MBPO on sparse-reward tasks, especially Hand, while preserving performance on dense-reward DMC tasks.Source: E MBPO · p. 17, Figure 18 · p. 17
  • MBPO with reward smoothing learns the first Hand task, whereas MBPO without smoothing fails on both Hand and RoboDesk.Source: E MBPO · p. 17

F CRAFTER

Crafter results show that smoothing kernels requiring future-reward anticipation perform worse, likely because they produce false positives that hinder policy learning.

F CRAFTER

  • Symmetric and future-only smoothing significantly deteriorate Crafter performance despite improving reward predictions.Source: F CRAFTER · p. 18
  • Symmetric kernels make reward predictions depend on past and future ground-truth rewards, requiring anticipation of future rewards.Source: F CRAFTER · p. 18
  • The asymmetric variants are Uniform [−δ, 0], which uses past rewards, and Uniform [0, δ], which uses future rewards.Source: F CRAFTER · p. 18
  • Symmetric and Uniform [0, 4] models produce more false positives than Uniform [−4, 0] and unsmoothed models.Source: F CRAFTER · p. 18, Figure 20 · p. 19, Figure 21 · p. 20
  • Uniform [−4, 0] performs best on Crafter, followed by symmetric uniform smoothing and then Uniform [0, 4].Source: F CRAFTER · p. 18, Figure 19 · p. 18

Related papers

Loading 2311.01450v2…