$λ$-models: Effective Decision-Aware Reinforcement Learning with Latent Models

TL;DR

This work investigates decision-aware model learning in model-based reinforcement learning, focusing on IterVAML and MuZero and the role of latent models. The authors prove that IterVAML with latent models can achieve an unbiased value function approximation in stochastic environments, while MuZero's value loss induces a bias under similar conditions; empirically, latent representations enhance performance and help stabilize training. They show that MuZero's performance gaps can be bridged by adopting latent models and stabilizing losses, and that decision-aware losses outperform BYOL in challenging, high-dimensional tasks such as Humanoid-run. The paper provides theoretical results, empirical validations, and practical guidance (the Lambda family) for when and how to apply decision-aware losses, especially in settings with limited model capacity or stochastic dynamics. Overall, it offers a nuanced view linking architectural choices, loss design, and environment dynamics to actionable recommendations for robust continuous-control RL.

Abstract

The idea of decision-aware model learning, that models should be accurate where it matters for decision-making, has gained prominence in model-based reinforcement learning. While promising theoretical results have been established, the empirical performance of algorithms leveraging a decision-aware loss has been lacking, especially in continuous control problems. In this paper, we present a study on the necessary components for decision-aware reinforcement learning models and we showcase design choices that enable well-performing algorithms. To this end, we provide a theoretical and empirical investigation into algorithmic ideas in the field. We highlight that empirical design decisions established in the MuZero line of works, most importantly the use of a latent model, are vital to achieving good performance for related algorithms. Furthermore, we show that the MuZero loss function is biased in stochastic environments and establish that this bias has practical consequences. Building on these findings, we present an overview of which decision-aware loss functions are best used in what empirical scenarios, providing actionable insights to practitioners in the field.

Figures (11)

• Figure 1: Sketch of the different value-aware losses with latent models. IterVAML computes the value function difference between the latent prediction and the next state encoding, while MuZero computes a single-step bootstrap estimate.
• Figure 2: Comparison on the use of an explicit latent space with different loss functions. All losses improve in performance from the addition of an explicit latent space transformation. Dashed lines represent the mean of the latent space result.
• Figure 3: The source of the MuZero bias. Though $x_1$ and $\hat{x}_1$ are drawn from the same distribution, their values do not coincide.
• Figure 4: Comparison of the impact of noise across all algorithms. While both IterVAML and MuZero are impacted by the level of noise, we observe a larger drop in MuZero, which does not perform above the BYOL baseline at the highest noise level.
• Figure 5: Performance comparison overall test environments. We see that IterVAML performs slightly above MuZero in several environments, but decisive gains from the value-aware losses ($\lambda$-IterVAML, $\lambda$-MuZero) over BYOL can only be observed in the challenging Humanoid-run environment.

Theorems & Definitions (8)

• Proposition 1
• Proposition 2
• Lemma 1: Deterministic Representation Lemma
• Lemma 2
• Proposition 1
• Proposition 2
• Proposition 3: Proposition 7.30 of BertsekasShreve1978
• Proposition 4: Proposition 7.33 of BertsekasShreve1978