DRIVE: Distributional Model-Based Reinforcement Learning via Variational Inference

Authors: Jianfei Ma, Huazhe Xu
Primary area: reinforcement learning
Keywords: Distributional RL, Control as Inference, Decision-Making under Uncertainty

Abstract:
Distributional reinforcement learning (RL) provides a natural framework for estimating the distribution of returns rather than a single expected value. However, the control aspect of distributional RL has not been as thoroughly explored as the evaluation part, typically relying on the greedy selection rule with respect to either the expected value, akin to standard approaches, or risk-sensitive measures derived from the return distribution. On the other hand, casting RL as a probabilistic inference problem allows for flexible control solutions utilizing a toolbox of approximate inference techniques; however, its connection to distributional RL remains underexplored. In this paper, we bridge this gap by proposing a variational approach for efficient policy search. Our method leverages the log-likelihood of optimality as a learning proxy, decoupling it from traditional value functions. This learning proxy incorporates aleatoric uncertainty of the return distribution, enabling risk-aware decision-making. We provide a theoretical analysis of our framework, detailing the conditions for convergence. Empirical results on vision-based tasks in DMControl Suite demonstrate the effectiveness of our approach compared to various algorithms, as well as its ability to balance exploration and exploitation at different training stages.

Decision: reject
Reviewer scores: 3, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730404726356

thread_type:
official_review

<thread_object>
creator: reviewer_emnd

<content>
summary:
The authors propose a framework that leverages probabilistic inferencing to find locally-optimal policy from distributional value function. In particular, they propose a variational approach that, amongst other things, decouples the optimization problem from the regular value function. They provide a theoretical analysis in support of this framework, as well as empirical results in a broad range of continuous control tasks.

strengths:
This paper provides a theoretical framework to find optimal policies in distribution RL, in which policy optimization does not fully utilize the distributional knowledge of value function in the existing literature. This is to complete the distributional approach to RL from end to end.

weaknesses:
The ideas presented in the paper have potential, but overall, it lacks motivational discussion (or justification) on various components in the proposed framework: why model-based approach by address H-horizon, and Assumption 3.1. Also, discussion is missing on the impact of approximate transition model, the additional burden from model-based approach, impact of H. 

The presentation of the experiments needs to be improved. In particular, by aggregating 11 experiments into a single figure (Fig. 2), we lose the ability to generate any meaningful insights as to which conditions may favor the author’s methods. Moreover, only some baselines are included in Figure 2, and soft-actor critic is excluded from Table 1, which is curious and not motivated in the text. Overall, there is no intuition provided by the authors as to why it makes sense to compare their method to the other methods, besides D4PG.

Finally, regarding the method proposed by the authors, I also have some concerns. First, the method appears to be much more computationally-expensive than, say, the more standard Distributional RL methods by Bellemare et al. (2017), due to the marginalization over the return, as well as the inclusion of Dreamer. Similarly, the inclusion of Dreamer would likely introduce some error into the learning processes, however this is not discussed or investigated by the authors.

questions:
1.	Would Eq (4) be compatible with the Bellman Eq (2). Wouldn’t Eq (4) lead to entropy term in the reward?
2.	How do you define U_max in Assumption 3.1?
3.	The transition associated with the second inequality in Eq (9) is not clear to me.
4.	Where does Eq (13) come from?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730637690554

thread_type:
official_review

<thread_object>
creator: reviewer_6hng

<content>
summary:
This paper proposes DRIVE, a novel distributional model-based reinforcement learning (RL) framework leveraging variational inference to address limitations in traditional RL approaches focused on expected return maximization. DRIVE introduces probabilistic learning proxies that incorporate aleatoric uncertainty within the return distribution, promoting risk-aware decision-making. The theoretical analysis demonstrates conditions for convergence, and empirical results on tasks from DMControl Suite highlight DRIVE's efficiency and its ability to balance exploration and exploitation.

strengths:
Originality:
The integration of variational inference into distributional RL is novel, and facilitates a risk-aware decision-making approach

Significance:
The paper is well theoretically-founded and appears to be effective

Clarity:
The paper is well-written and logically structured

weaknesses:
- Although the paper emphasizes balancing exploration and exploitation, it lacks a detailed analysis of how can DRIVE balance exploration and exploitation during traning, it would be helpful to include a detailed discussion.

- The multi-term objective and the complex posterior approximation might lead to overfitting, especially in limited-data scenarios. Specifically, how does your method scale to high-dimensional and large action-space environments

- Missing reference: The paper does not address a comparison with bisimulation mode based approaches [1, 2, 3]. How does your method relate to representation learning [1] and reward shaping methods [2,3] that also utilize the bisimulation metric for learning transition and reward models? These methods can be categorized as model-based approaches connected to the optimal value function. Specifically, [3] also manages to encourage exploration without loss of exploitation by connecting the shaping reward with the value function. What advantages does DRIVE offer over these bisimulation model-based approaches?

[1] Zhang et al. "Learning invariant representations for reinforcement learning without reconstruction." ICLR 2021

[2] Kemertas et al. "Towards robust bisimulation metric learning." NeurIPS 2021

[3] Wang et al. "Efficient potential-based exploration in reinforcement learning using inverse dynamic bisimulation metric." NeurIPS 2023

questions:
- Can the authors provide more concrete evidence or metrics that demonstrate how their method manages the exploration-exploitation trade-off across different tasks?

- How does the complexity of the variational bound affect the scalability of the method in large state-action spaces?

- What measures, if any, have been implemented to mitigate overfitting due to the complex objective function and posterior approximation
Speficifically, how does DRIVE perform in discrete-action games, e.g., atari games?

- What advantages does DRIVE offer over these bisimulation model-based approaches[1,2]? (see W3)

- Some parts of the theoretical analysis are dense. Could you provide more intuitive explanations or examples to help bridge the gap between the theory and its practical implications?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730715334863

thread_type:
official_review

<thread_object>
creator: reviewer_h7os

<content>
summary:
The paper combines perspectives from RL as inference and distributional RL nicely through introducing the return random variable $U$, which offers a unique perspective as well as a new principled way to do policy search. The perspectives introduced in this paper can be helpful for future works that combine tools from RL, distributional RL, and probabilistic inference. Empirical results are reasonable.

strengths:
**[Originality]**
The paper combines perspectives from RL as inference and distributional RL nicely through introducing the return random variable $U$, which offers a unique perspective as well as a new principled way to do policy search. 

**[Quality]**
The idea is well motivated and empirical results are reasonable. 

**[Clarity]**
The writing and derivations are for the most part easy to follow

**[Significance]**
The perspectives introduced in this paper can be helpful for future works that combine tools from RL, distributional RL, and probabilistic inference.

weaknesses:
For me, the main weaknesses are incomplete empirical evaluations, and some (for me) confusing choices made in the variational lower bound derivations. 

- In Figure 4, the 95 confidence interval (?) between ablations are overlapped, suggesting there is no statistically significant difference between their mean. More seeds can be helpful to separate out the differences. 
- Taking more samples to approximate a distribution ($\mathfrak{q}$, see L248) actually results in *worse* performance (Figure 4b) is strange, and in my opinion the authors do not provide a sufficiently satisfactory explanation for this. 
- It is not fully clear yet how the method’s behavior differs from existing work. The method seems to be roughly on-par with existing strong model-based methods (Figure 2), with potentially interesting differences in their exploration (Figure 3) that can be further investigated. 
- I will ask about variational approximation in the “question” section below.

questions:
1. I am confused by equation 29 (which impacts the lower bound derivation of equation 9). Equation 29 can be equivalently written as 
$$\log p_\psi (\mathcal{O}=1 | s,a) = \log E_{p_\psi(U|s,a)} [\frac{\exp U}{ \exp U_{max}} ] \geq E_{p_\psi(U|s,a)}[U] - U_{max}$$ 
without the need to introduce distribution $\mathfrak{q}$. Since $p_\psi(U|s,a)$ is supposed to be the parametric return model (L183) which we can sample from (or even have in close form), why did the authors introduce an additional distribution $\mathfrak{q}$ (in Eq 30)? 

2. Relating to above, $\mathfrak{q}$ seems to be used as the “true” return distribution (e.g. Eq 11), which makes it weird that it does not “need” to appear in the lower bound at all. Am I conceptually missing something here? It perhaps feels like $\log p^\pi_{\psi} (\mathcal{O}=1 | s)$ is not quite the correct quantity to lower-bound and optimize (but should instead be e.g. a quantity that does not depend on $\psi$)?

3. What is the parametric distribution $p_\psi(U|s,a)$, I presume it is a Gaussian? 

4. For results: why are the per-game SAC and DreamerV3 results from Figure 2 not present in Table 1? Also, can the authors provide per-game performance for the 11 games for SAC and DreamerV3?

5. Clarity: for exploration, can the authors provide more details about the two criterias being compared in Figure 3, bottom (i.e. in equations)? Further, I do not feel Figure 3 top is sufficiently well-explained to make sense. 

6. The effect of this new objective on exploration is very interesting, and I think worth investigating further. For instance by quantifying state coverage, and/or showing an illustrative exploration trace in simple environments? As an example, see Figure 2 and 7 of [Li 2024]. 

7. For clarity, can the authors provide equations of the objectives when “replacing the posterior with the policy and removing the regularizer term”? L451.

8. (L448) The word “disentanglement” has a lot of associated meanings in machine learning. Perhaps just “ablation” would suffice as the paragraph title. 

9. Nit pick: please use “\citep” and \citet” properly. For eg citations on L40, L57 (amongst others) should only be in a single bracket, and citations such as Dabney et al on L488 should use “citet”. 

[Li 2024] Li, Qiyang, et al. "Accelerating exploration with unlabeled prior data." Advances in Neural Information Processing Systems 36 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730803489719

thread_type:
official_review

<thread_object>
creator: reviewer_m9wt

<content>
summary:
This paper proposes a combination of distributional RL and probabilitstic inference in model based RL to incorporate aleatoric uncertaity of distribution for risk-aware decision-making.

The main learning objecitve of this paper is the log-likelihood of the optimality variable that is computed by marginalizing the conditional probability $p(O=1|U,s,a)$ on $U$ where $U$ is the return distribution.

Since the optimality variable $O=1$ implies the optimal return at given state and action, maximizing the return of agent induces the log-likelihood of optimality variable.

Compared to Dreamer, the above log-likelihood is added to the learning objective.

In experiments, all experiments is conducted in vision-based deepmind control suite.

strengths:
The main contribution of this paper is to reduce the effect of aleotoric uncertainty in training model-based RL.

By introducing the optimality variable, the proposed method can asses the quality of return distribution by imagination and match to the ground-truth return distribution better by preventing from leaning to the maximum value.

weaknesses:
The main conern that I have is the quality of presentation and the lack of experiments.

1. The notations are not consistent. For example, there are several different q functions which outputs  $a$ in Equation 9, $s_t$ in Equation 16, $U$ in Algorithm 1 without any proper subscriptions.
2. The theoretical results, theorem 5.1 and theorm 5.2 are not connected to the main contribution which induces the better performance of model-based distributional RL method.
3. It is not clear which distributinoal RL method is used to estimate $U$, such as C51 or QRDQN.
4. The experiments is only conducted in deepmind control suite, which is not reported in the main baseline, Dreamer. If the proposed method has an advantage in vision-based continuous control, the authors should have described how to perform better in continous control tasks.

questions:
Can you make the notations be clear to help to understand the main contribution better?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732630529213

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Changes

comment:
We greatly appreciate the constructive comments from the reviewers. In response, we have revised the paper based on their suggestions. The updates are summarized as follows:

* We provided further evidence of balancing exploration and exploitation by designing a minimal, diagnostic multi-armed bandit experiment at the request of **`Reviewer h7oS`** and **`Reviewer 6hNG`**.
* We added a comparison of second per iteration in **Table 4(c)** to address the computational concern raised by **`Reviewer EmnD`**.
* We provided the full task-wise results for all baselines in **Table 1** at the request of **`Reviewer h7oS`** and **`Reviewer EmnD`**.
* We added a visualization and textual explanation in Figure 2 to aid understanding.
* We revised the font style for the q function in both Eq (15) and (16), and added subscripts for $U$ in **Algorithm 1** to enhance clarity as pointed out by **`Reviewer m9WT`**.
* We provided additional experimental results in **Appendix C.4**: one monitoring the policy entropy and the other investigating the impact of the planning horizon, as requested by **`Reviewer EmnD`**.

We have highlighted all the revised/added parts in **dark blue** for the convenience of the reviewers. We hope those changes address the reviewers' concerns, and we look forward to receiving their feedback.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a new model-based distributional RL technique based on variational inference.  It claims to enable risk sensitive decision making.  However this claim is not supported by suitable theoretical or empirical evidence.  The paper observes that previous techniques in control as inference define the probability of optimality to be proportional to the expected return of a trajectory.  Similarly, the paper proposes to use the exponential of the sampled return of a trajectory as a surrogate of optimality.  This is an interesting idea, but it lacks a clear mathematical derivation that demonstrates how risk sensitive decision making is achieved.  In fact, since the paper proposes to define the probability of optimality in terms of exponential returns, it seems that the literature on transformations of returns applies here.  More specifically, transforming returns with a concave function enables risk averse decision making while transforming returns with a convex function enables risk seeking decision making.  Since the paper proposes to take the exponential of the returns, it seems to enable risk seeking decision making only (not risk averse decision making).  See https://en.wikipedia.org/wiki/Risk_aversion for a discussion.  This needs to be clarified.

Ultimately, the theory presented simply analyzes the convergence of the algorithm, but does not show that the resulting policy can be risk averse nor how to adjust the degree of risk aversion.  Similarly, the experiments do not clearly demonstrate how to obtain risk averse policies nor how to adjust the degree of risk aversion.

While the paper investigates an interesting construction, there is a lack of justification and evidence for the proposed construction.  Hence this work is not ready for publication.