Direct Advantage Estimation in Partially Observable Environments

Authors: Hsiao-Ru Pan, Bernhard Schölkopf
Primary area: reinforcement learning
Keywords: POMDP, advantage function, deep RL, off-policy learning

Abstract:
Direct Advantage Estimation (DAE) was recently shown to improve sample-efficiency of deep reinforcement learning algorithms. However, DAE assumes full observability of the environment, which may be restrictive in realistic settings. In the present work, we first show that DAE can be extended to partially observable domains with minor modifications. Secondly, we address the increased computational cost due to the need to approximate the transition probabilities through the use of discrete latent dynamics models. Finally, we empirically evaluate the proposed method using the Arcade Learning Environments, and show that it is scalable and sample-efficient.

Decision: reject
Reviewer scores: 6, 6, 5, 5

Reviews and discussion:
<thread>
creation date:
1729810682546

thread_type:
official_review

<thread_object>
creator: reviewer_jrnv

<content>
summary:
This paper extend the off-policy Direct Advantage Estimation (DAE) technique into the POMDP environments with model estimation trick to reduce the complexity.

strengths:
The idea is sound and motivation is intuitive with straightforward explanation, building upon the previous works of DAE and off-policy DAE[4,5]. The method itself has been proved effective on the Arcade Learning benchmarks. Authors have also conducted extensive experiments to validate hyper-parameters selection and various corrections.

weaknesses:
1, line 247, 255, and 266 would be better to refer to Eq 9 instead of Eq 14 back in the appendix.\
2, Though technical sound, this paper seems a bit incremental comparing with previous works[4] with a different dynamic reparameterization for high-dim state space.\
3, Section 3.2 states the reduction of computational complexity comparing to LSTM and RNN based methods, but it would be great to have some empirical or theoretical demonstrations.\
4, The authors claim that off-policy correction is crucial for stochastic environment and performance improvement. It would be better to explain how the stochastic part is validated. \
5, Reviewer is also interested in seeing the comparison with other recent works with support for pomdp such as [1,2,3]. Current comparisons with Rainbow and DQN might not be sufficient. 

[1] Hansen, Nicklas, Hao Su, and Xiaolong Wang. "Td-mpc2: Scalable, robust world models for continuous control." arXiv preprint arXiv:2310.16828 (2023).\
[2] Chen, Xiaoyu, et al. "Flow-based recurrent belief state learning for pomdps." International Conference on Machine Learning. PMLR, 2022.\
[3] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pages 2555–2565. PMLR, 2019.\
[4]Pan, Hsiao-Ru, and Bernhard Schölkopf. "Skill or luck? return decomposition via advantage functions." arXiv preprint arXiv:2402.12874 (2024).\
[5]Pan, Hsiao-Ru, et al. "Direct advantage estimation." Advances in Neural Information Processing Systems 35 (2022): 11869-11880.

questions:
1, Is it necessary to incorporate reward signal $r_{0:t-1}$ to the information vector $h_t$? Since $r_t$ info seems not be incorporated in the $h_t$.\
2, The reviewer is a little bit confusing on the mixing use of state $s$, $h$, $z$, and $x$. It would be great to have a better clarification on these symbols and definitions of corresponding $\hat{A}, \hat{B},$ and $\hat{V}$.\
</content>
</thread_object>

</thread>
<thread>
creation date:
1730363213264

thread_type:
official_review

<thread_object>
creator: reviewer_pvss

<content>
summary:
This paper presents an extension of the off-policy Direct Advantage Estimation (DAE) method to partially observable Markov Decision Processes (POMDPs). The proposed approach demonstrates improved performance on certain benchmark problems compared to DQN. However, the notation throughout the paper is unclear, making it difficult to follow the methodology. Additionally, the baseline algorithm used for comparison is outdated, limiting the persuasiveness of the results. The motivation for applying DAE in this context is also insufficiently explained, and I find the necessity of the proposed method unconvincing.

strengths:
1. This paper extend off-policy DAE to POMDP with minor modifications.
2. Authors address the problem of increasing computational costs by modeling the latent space transition with RNN.
3. Partial experimental results demonstrate that the proposed method has good performances.

weaknesses:
1. The paper is hard for me to follow, and the writing could benefit from further polishing. For instance, the parameter notation $\theta$ is used in various forms, such as $f_\theta(h, a)$, $g_\theta(s, a, z)$, and $p_\theta(z|h_t, a_t)$. This raises questions about whether these functions share the same parameters, which the paper does not clarify. Additionally, the notation for $z$ in Eq. (10) is undefined until $\mathcal{Z}$ is introduced before Eq. (11), making the notation confusing.

2. The authors appear to consider only the simplest case, with finite state and action spaces. Under these conditions, why is DAE necessary, and why wouldn’t a TD loss with a Bellman operator suffice? The motivation for introducing DAE in this context is unclear and could be better presented.

3. The benchmark algorithms used in this paper date back 6-9 years, which weakens the strength of the comparison. I recommend that the authors incorporate more recent, robust baselines to more convincingly showcase the advantages of their proposed method.

questions:
Same as the weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730667398828

thread_type:
official_review

<thread_object>
creator: reviewer_6qrc

<content>
summary:
The paper extends and improves upon Off-Policy Direct Advantage Estimation, both by extending it to the POMDP case by modifying the way the loss is computed, but also by modelling the transitions in the latent space of a VQ-VAE, which improves the computational cost of training models with DAE.

strengths:
The theoretical analysis is quite complete and sound, and the authors bring techniques from other subfields of Reinforcement Learning (World Models) to improve DAE and make it more practical.

weaknesses:
The experimental analysis is rather weak, they compare their model with Rainbow and DQN, while using a recurrent and residual architecture signifcantly more complex and costly. Comparing with Impala and R2D2 would have been more meaningful. Besides that they report only mean and 1 standard error of their agents, whereas reporting the median agent performance and either a 2 standard deviation interval or better yet a tolerance interval would have improved the visualization of their algorithm performance.

questions:
1. Would it be possible to compare the performance with methods using similar networks to the ones you used, like Impala and R2D2? Why did you think this comparison was irrelevant?
2. Would it be possible to report ablations on using a Exponential Moving Average network vs a periodically updated network (DQN-style) as your target network? If not would it be possible to justify the choice a bit more?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691001674

thread_type:
official_review

<thread_object>
creator: reviewer_nsq7

<content>
summary:
The paper extends and improves an existing off-policy reinforcement learning (RL) algorithm called direct advantage estimation (DAE) to the setting of partially observable environments. 

The paper motivates learning advantage functions, as opposed to value functions, given their more stable learning dynamics. Then, it explains how DAE can work in the off-policy case by separating the return in two advantage functions: one associated with the actions of the agent and the other with the transitions, interpreted as environment actions. Based on this separation, the paper re-introduces value and advantage functions as functions of observation (and action and reward) histories, rather than states, and provides an optimization problem whose minima contain the real value and advantage functions. 

Moreover, it is explained that, since some of the constraints in the optimization problem depend on having access to the real dynamics, there is a need for some type of approximate model. As a solution, the paper introduces a predictive recurrent architecture that models the dynamics of the environment in an embedding space, mimicking a discrete version of variational autoencoders (VQ-VAEs). 

Finally, the paper evaluates the performance of the algorithm based on return maximization in 5 different Atari environments, including some ablations. The results show that the algorithm can compete with state-of-the-art methods like Rainbow with a considerably smaller number of samples, that the off-policy correction introduced by the use of the environment's advantage function is critical to compete with existing techniques, that there are diminishing returns on the multi-step size considered to estimate advantages, that frame-stacking can be replaced by an LSTM given that the algorithm supports partial observability, and that naive truncation of the history can be detrimental.

strengths:
Originality:
- The paper makes use of a simple modification to a previously introduced concept to be able to leverage previous theory. In particular, the paper redefines the environment's advantage function in terms of the reward and V- and Q- value functions, as opposed to just the V-value function. This allows transferring an existing algorithm to the more challenging partially observable setting.
- The paper provides a conceptually simple predictive model in encoding space to manage the approximation of the dynamics.

Quality:
The paper is thorough in its empirical evaluation. In particular, it includes a sufficient number of seeds, relevant baselines, and multiple ablations. Some of the ablations are particularly appropriate since they provide some evidence that the algorithm does maximize return under partial observability.

Clarity:
In general, the paper is easy to read due to its simplicity and clarity. It introduces very succinctly the problem and clearly states the contributions. It also introduces all the relevant notation and explains in detail the experiments carried out.

Significance:
The paper provides convincing evidence to conclude that the proposed algorithm is state-of-the-art. In particular, it is shown that the algorithm makes use of fewer samples to obtain similar or better performance than the baselines in medium size environments that require the use of function approximation and handling some degree of partial observability.

weaknesses:
In my opinion, there is only one weakness worth focusing on. The main contribution is supposed to be the extension of DAE to the partially observable setting, but the Atari environments chosen are arguably not challenging from the partially observable perspective. In particular, the fact that DQN and Rainbow, algorithms that are typically paired with non-recurrent architectures, can maximize rewards in them suggests that partially observability is not a big issue on them. In this sense, I think that ablation studies in these environments, for example, one comparing frame-stacking vs. LSTMs, are not as effective as choosing tasks where partial observability is a crucial factor for good performance. 

Now, while I consider that the paper should be accepted in its current version, I think that a higher score requires considering more appropriate environments. For example, see Pašukonis et al., Evaluating Long-Term Memory in 3D Mazes, in ICLR, 2023.

questions:
- What is the computational overhead of the predictive model?
- In line 246, why is $\hat{x}$ being called an information vector?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732038371436

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response

comment:
We thank the reviewers for the valuable feedback and thoughtful comments, which have greatly helped improve the quality of this work. Here, we summarize the changes:

1. **Baselines**: As pointed out by the reviewers, the baselines (DQN/Rainbow) used in this paper, while standard for the ALE, are no longer considered SOTA anymore. Consequently, we have updated our list of baselines to include more recent SOTA model-based recurrent algorithms DreamerV2 [1] and DreamerV3 [2]. 
2. **Notations**: As noted by the reviewers, there were some inconsistencies in the usage of state variables ($s$) and information vectors ($h$), which could be confusing. To address this, we have updated the notations in Section 3 to ensure better consistency across different functions. Additionally, we have revised Figure 1 to more clearly illustrate the connections between the variables ($o_t, x_t, h_t$) and the VQ-VAE.
3. **New ablation**: As suggested by reviewer 6QRc, we have added a new ablation study (see Appendix C.5) that compares periodic updates of the target, common in DQN variants, to EMA-style updates, more typical in self-supervised methods like SPR.

We thank the reviewers once again and hope that our responses, along with the revised manuscript, address any confusion and resolve all outstanding issues.

[1] Hafner, Danijar, et al. "Mastering atari with discrete world models." arXiv preprint arXiv:2010.02193 (2020).

[2] Hafner, Danijar, et al. "Mastering diverse domains through world models." arXiv preprint arXiv:2301.04104 (2023).
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper extends prior work on off-policy Direct Advantage Estimation (DAE) by enabling support for partially-observable systems. The authors provide theoretical justification for their approach and couple it with an empirical evaluation on the ALE.

All reviewers agree that the paper is clearly written (especially after revisions during the rebuttal), the method interesting, and the theory well presented.

However, given how heavily the proposed method builds on DAE with the claim that this enables running on partially observable environments, it was felt that the provided empirical evidence is not sufficient to fully justify its need. In particular, all reviewers agreed that more games on the ALE would be useful. Although the five were chosen based on the suggestion of the Atari-5 paper, it is not clear that those 5 are the best for highlighting issues with partial observability, which is the main focus of this work.

Perhaps showcasing this method in a setting where the distinction between full and partial observability (beyond frame-stacking) would clarify the distinction.

In summary, this is interesting work that unfortunately does not quite yet meet the bar for publication, for the reasons mentioned above.